)\n```\n\n### Examples\n\n```\nreshape: cannot reshape [10,10] to [12,10] (100→120 elements)\n\nbroadcast: cannot broadcast [2,3] with [4,5] (dim 0: 2≠4, dim 1: 3≠5)\nhint: broadcasting requires dimensions to be either equal or 1\n\nempty: invalid shape [-1, 10] (negative dimension)\n\nmatmul: cannot multiply Float32 @ Int64 (dtype mismatch)\nhint: cast one array to match the other's dtype\n```\n\n### Rules\n\n#### Always include:\n- **Operation name** - what function failed\n- **Full context** - complete shapes, not just sizes\n- **Specific problem** - which dimension/axis failed and why\n\n#### Structure consistently:\n- For transformations: `[10,10] to [12,10]`\n- For operations: `[2,3] with [4,5]`\n- For access: `[5,2] in shape [3,4]`\n- For invalid inputs: `invalid X (reason)`\n\n#### Make problems obvious:\n- Show comparisons: `2≠4`, `5≥3`, `100→120`\n- Point to location: `dim 0:`, `axis 1:`\n- State violations: `axis 2 repeated`, `multiple -1`\n\n#### Multiple issues:\n```\nconv2d: invalid configuration\n - input channels: 3 ≠ 5 (weight expects 5)\n - kernel [6,6] > input [5,5] with 'valid' padding\n```\n\n#### Add hints when:\n- The fix is non-obvious\n- There's a specific function to call\n- The rule isn't clear from context\n- Backend limitations exist\n\n### Special Cases\n\n**Performance warnings:**\n```\nreshape: requires copy from strided view [100,10] to [1000]\nhint: call contiguous() first to avoid copy\n```\n\n**Empty/scalar edge cases:**\n```\nsqueeze: cannot squeeze scalar (already rank 0)\nargmax: empty axis returns no indices (size 0)\n```\n\n**Backend limitations:**\n```\ngather: indices dtype Int64 not supported (backend uses Int32)\nhint: cast indices to Int32\n```\n\n### Common Patterns\n\n**Shape changes:**\n```\nreshape: cannot reshape [2,5,10] to [4,26] (100→104 elements)\n```\n\n**Invalid access:**\n```\nslice: cannot slice [(0,5), (2,12)] in shape [10,10] (axis 1: 12>10)\n```\n\n**Type/value errors:**\n```\npad: invalid padding [-1, 2] (negative values)\nhint: use shrink() to remove elements\n```\n\n**Configuration errors:**\n```\npermute: invalid axes [0,2,2] (axis 2 repeated)\narange: invalid range [10, 5, 1] (start > stop with positive step)\n```\n\n### Don'ts\n\n❌ Vague errors: `invalid shape`\n❌ Missing context: `100 != 120`\n❌ Redundant hints: `shapes must be compatible (incompatible shapes)`\n❌ Teaching basics: `broadcasting requires...` (save for hints)\n\n### Summary\n\nShow exactly what they tried, what failed, and where. Use the standard format when possible, adapt when needed. Include hints only when they add value.\n"
},
{
"path": "LICENSE",
"content": "ISC License\n\nCopyright (c) 2025, Thibaut Mattio\n\nPermission to use, copy, modify, and/or distribute this software for any\npurpose with or without fee is hereby granted, provided that the above\ncopyright notice and this permission notice appear in all copies.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\" AND THE AUTHOR DISCLAIMS ALL WARRANTIES\nWITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF\nMERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR\nANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES\nWHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN\nACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF\nOR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE."
},
{
"path": "README.md",
"content": "\n
\n\nmodern scientific computing for OCaml \n\n\n docs ·\n install ·\n issues \n
\n\n---\n\nRaven is an ecosystem of OCaml libraries for numerical computing, machine learning, and data science. Everything you know from Python — NumPy, JAX, PyTorch, Matplotlib, Jupyter — rebuilt with type safety.\n\n> Raven is **alpha**. APIs will change. [Feedback welcome.](https://github.com/raven-ml/raven/issues)\n\n```ocaml\n(* nx — n-dimensional arrays *)\nlet x = Nx.linspace float32 0. 10. 100\nlet y = Nx.sin x\n\n(* rune — automatic differentiation *)\nlet grad_f = Rune.grad (fun x -> Rune.sum (Rune.mul x x)) x\n\n(* brot — tokenization *)\nlet tokenizer = Brot.from_file \"tokenizer.json\" |> Result.get_ok\nlet ids = Brot.encode_ids tokenizer \"The meaning of life is\"\n\n(* kaun — neural networks *)\nlet model = Kaun.Layer.sequential [\n Kaun.Layer.linear ~in_features:768 ~out_features:128 ();\n Kaun.Layer.relu ();\n Kaun.Layer.linear ~in_features:128 ~out_features:10 ();\n]\n\n(* talon — dataframes *)\nlet df = Talon.create [\n \"name\", Talon.Col.string_list [ \"Alice\"; \"Bob\"; \"Charlie\" ];\n \"score\", Talon.Col.float64_list [ 85.5; 92.0; 78.5 ];\n]\n\n(* hugin — plotting *)\nlet () = Hugin.(figure () |> subplot |> Plotting.plot ~x ~y |> ignore; show ())\n```\n\n## Packages\n\n| | Package | Like | What it does |\n| --- | ------------------------------ | ----------------- | -------------------------------------------------------- |\n| | [**nx**](packages/nx/) | NumPy | N-dimensional arrays with linear algebra operations |\n| ᛏ | [**tolk**](packages/tolk/) | tinygrad | Minimal ML compiler for GPU tensor computation |\n| ᚱ | [**rune**](packages/rune/) | JAX | Automatic differentiation and functional transformations |\n| ᚲ | [**kaun**](packages/kaun/) | Flax | Neural networks and training |\n| ᚹ | [**vega**](packages/vega/) | Optax | Composable gradient-based optimizers |\n| ᚾ | [**norn**](packages/norn/) | BlackJAX | MCMC sampling with automatic gradients |\n| ᚨ | [**brot**](packages/brot/) | HF Tokenizers | Fast, HuggingFace-compatible tokenization |\n| ᛃ | [**talon**](packages/talon/) | Polars | Fast and elegant dataframes with type-safe operations |\n| ᛞ | [**hugin**](packages/hugin/) | Matplotlib | Publication-quality plotting |\n| ᛈ | [**quill**](packages/quill/) | Jupyter + IPython | Interactive REPL and markdown notebooks |\n| ᚠ | [**fehu**](packages/fehu/) | Gymnasium | Reinforcement learning environments |\n| ᛋ | [**sowilo**](packages/sowilo/) | OpenCV | Differentiable computer vision |\n| ᛗ | [**munin**](packages/munin/) | W&B / MLFlow | Local experiment tracking with live TUI dashboard |\n\n## Getting started\n\n```bash\nopam install raven\n```\n\nThis installs the full ecosystem. You can also install only what you need — e.g. `opam install kaun` for neural networks, or `opam install nx` for just arrays.\n\nAdd to your `dune` file:\n\n```dune\n(executable\n (name main)\n (libraries raven))\n```\n\nSee the [installation guide](https://raven-ml.dev/docs/installation/) for system dependencies and editor setup.\n\n## Support\n\nBuilding a scientific computing ecosystem takes sustained effort. Sponsorships help us ship JIT compilation, distributed training, better developer tooling, and production deployment through MirageOS.\n\n**[Support Raven →](https://raven-ml.dev/docs/support-raven/)**\n\nThanks to our sponsors [Ahrefs](https://ahrefs.com) and [Tarides](https://tarides.com).\n\n## License\n\n[ISC](LICENSE)\n"
},
{
"path": "TODO.md",
"content": "# todo\n\n## beta (jit)\n\ngoalpost: jit-compiled gpt2 matching pytorch performance\n\nperf:\n- close rune grad performance gap (within <2x of pytorch)\n- close nx performance gaps (within <2x of numpy)\n\ntolk:\n- integrate tolk as rune jit transformation\n- kernel fusion and optimization\n- cpu, cuda, metal backends\n\n## v1 (production)\n\ngoalpost: end-to-end train -> deploy as unikernel or static binary\n\ntraining:\n- gradient accumulation\n- mixed precision (fp16/bf16 forward, fp32 master weights, loss scaling)\n- gradient checkpointing (rune.checkpoint, recompute activations in backward)\n- flash attention (tolk kernel and/or kaun.fn primitive)\n- parallel data loading (ocaml 5 domains, background prefetch)\n- layer completions: transposed conv, group norm, full conv2d stride/dilation/padding\n- onnx import (onnx -> tolk ir adapter, cover resnet/bert/gpt2/llama/vit/whisper ops)\n\ndeployment:\n- aot compilation: cpu (c via clang, musl static linking) and gpu (cuda/metal/opencl)\n- mimir: kv cache, continuous batching, pagedattention\n- mimir: http server (rest api, /health, /metrics, sigterm, structured logging)\n- post-training quantization (int8/int4, tolk quantized kernels)\n- mirageos unikernel deployment (raven-mirage package)\n - no blas dep (tolk aot generates all compute)\n - weight loading via network (mirage-http)\n - verify ocaml 5 effects on mirageos runtime\n - http server on mirageos network stack\n\ndocs/website:\n- landing page rewrite with benchmarks\n- deployment guide (aot, static binary, docker, mirageos, gpu)\n- end-to-end examples (serving, onnx+deploy workflow)\n"
},
{
"path": "dev/README.md",
"content": "# dev\n\nDevelopment sandbox for experiments and prototypes that support the Raven ecosystem.\n\n## Projects\n\n| Name | Description |\n| ---- | ----------- |\n| [mimir](mimir/) | Experimental inference engine |\n| [tolk](tolk/) | ML compiler inspired by tinygrad |\n"
},
{
"path": "dev/mimir/README.md",
"content": "# mimir\n\nExperimental inference engine for raven.\n\nThe gap between \"I can run a forward pass\" and \"I can serve a model in production\" is large. mimir is where we figure out what the OCaml answer to that gap looks like.\n\n## Current state\n\nThe sampling layer: composable logits processors (temperature, top-k, top-p, repetition penalty, n-gram blocking), stopping criteria, and the autoregressive generation loop operating on nx tensors.\n\nThis is the outermost piece of the inference puzzle — the part that turns model logits into actual token sequences. Everything below is open.\n\n## What we want to explore\n\n**Memory management for KV cache.** The attention mechanism produces intermediate state (keys and values) that grows linearly with sequence length. Naive allocation wastes memory; the interesting question is whether we can apply OS-style virtual memory ideas — fixed-size blocks, deferred allocation, reference-counted sharing — to make long sequences and shared prefixes cheap. This is the core idea behind PagedAttention.\n\n**Request scheduling.** A single request is simple. Thousands of concurrent requests with different prompt lengths, generation limits, and priority levels is a scheduling problem. Batching amortizes GPU overhead but introduces latency trade-offs. Continuous batching (letting new requests join mid-batch as others finish) changes the calculus further. OCaml's algebraic types and pattern matching may give us a cleaner expression of scheduling policies than the typical mutable-state approach.\n\n**Prefill/decode asymmetry.** The two phases of autoregressive generation have opposite performance characteristics — one is compute-bound, the other memory-bound. An engine that treats them identically leaves performance on the table.\n\n**JIT compilation of decode steps.** The decode phase repeats the same computation graph with different inputs. If rune's JIT can capture and replay these graphs, we avoid per-step compilation overhead — similar in spirit to CUDA graph capture.\n\n**Structured generation.** Constraining the sampling step so that output conforms to a grammar, regex, or JSON schema. This means masking logits at each step based on what the constraint automaton allows, which interacts with the sampling pipeline we already have.\n\n**Tensor parallelism.** Splitting a model across multiple devices. This is a rune-level concern more than a mimir concern, but the inference engine needs to coordinate it.\n\n## References\n\n- [Nano-vLLM](https://github.com/GeeeekExplworker/nano-vllm) — minimal (~1,200 lines) inference engine by a DeepSeek contributor, good for understanding the essential moving parts\n- [vLLM: PagedAttention paper](https://arxiv.org/abs/2309.06180)\n- [SGLang](https://github.com/sgl-project/sglang) — alternative engine with RadixAttention for prefix sharing\n"
},
{
"path": "dev/mimir/dune-project",
"content": "(lang dune 3.21)\n\n(name mimir)\n\n(package\n (name mimir)\n (synopsis \"Experimental inference engine for Raven\")\n (description\n \"Mimir is an inference engine for the Raven ecosystem. It provides sampling, KV cache management, request scheduling, and structured generation for serving ML models.\")\n (depends\n (ocaml\n (>= 5.2))))\n"
},
{
"path": "dev/mimir/lib/dune",
"content": "(library\n (name mimir)\n (public_name mimir)\n (libraries nx unix))\n"
},
{
"path": "dev/mimir/lib/mimir.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ninclude Sampler\n"
},
{
"path": "dev/mimir/lib/mimir.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Mimir - Text generation with composable logits processors.\n\n Experimental inference/generation library for the Raven ML ecosystem.\n Provides the autoregressive decode loop, composable logits processors,\n stopping criteria, and generation configuration. *)\n\ninclude module type of Sampler\n"
},
{
"path": "dev/mimir/lib/sampler.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* ───── Core Types ───── *)\n\ntype logits = (float, Bigarray.float32_elt) Nx.t\ntype token_ids = int array\n\n(* ───── Logits Processors ───── *)\n\ntype logits_processor = {\n name : string;\n process : prompt_length:int -> token_ids -> logits -> logits;\n}\n\ntype logits_processor_list = logits_processor list\n\n(* ───── Stopping Criteria ───── *)\n\ntype stopping_criterion = {\n name : string;\n should_stop : prompt_length:int -> start_time:float -> token_ids -> bool;\n}\n\ntype stopping_criteria_list = stopping_criterion list\n\n(* ───── Generation Configuration ───── *)\n\ntype generation_config = {\n max_length : int;\n max_new_tokens : int option;\n min_length : int;\n min_new_tokens : int;\n do_sample : bool;\n temperature : float;\n top_k : int;\n top_p : float;\n repetition_penalty : float;\n no_repeat_ngram_size : int;\n bad_words_ids : int list list;\n force_words_ids : int list list;\n pad_token_id : int option;\n bos_token_id : int option;\n eos_token_id : int option;\n eos_token_ids : int list;\n}\n\nlet default =\n {\n max_length = 100;\n max_new_tokens = None;\n min_length = 0;\n min_new_tokens = 0;\n do_sample = false;\n temperature = 1.0;\n top_k = 0;\n top_p = 1.0;\n repetition_penalty = 1.0;\n no_repeat_ngram_size = 0;\n bad_words_ids = [];\n force_words_ids = [];\n pad_token_id = None;\n bos_token_id = None;\n eos_token_id = None;\n eos_token_ids = [];\n }\n\n(* ───── Builder Pattern ───── *)\n\nlet with_temperature temperature config = { config with temperature }\nlet with_top_k top_k config = { config with top_k }\nlet with_top_p top_p config = { config with top_p }\n\nlet with_repetition_penalty repetition_penalty config =\n { config with repetition_penalty }\n\nlet with_max_length max_length config = { config with max_length }\n\nlet with_max_new_tokens max_new_tokens config =\n { config with max_new_tokens = Some max_new_tokens }\n\nlet with_min_length min_length config = { config with min_length }\nlet with_min_new_tokens min_new_tokens config = { config with min_new_tokens }\n\nlet with_no_repeat_ngram no_repeat_ngram_size config =\n { config with no_repeat_ngram_size }\n\nlet with_do_sample do_sample config = { config with do_sample }\n\n(* ───── Preset Configurations ───── *)\n\nlet creative_writing =\n default |> with_do_sample true |> with_temperature 0.8 |> with_top_p 0.9\n |> with_repetition_penalty 1.2\n |> with_no_repeat_ngram 3 |> with_max_new_tokens 512\n\nlet chat =\n default |> with_do_sample true |> with_temperature 0.7 |> with_top_p 0.95\n |> with_repetition_penalty 1.1\n |> with_max_new_tokens 512\n\nlet code_generation =\n default |> with_do_sample true |> with_temperature 0.2 |> with_top_k 5\n |> with_repetition_penalty 1.0\n |> with_max_new_tokens 1024\n\nlet factual =\n default |> with_do_sample true |> with_temperature 0.3 |> with_top_k 10\n |> with_repetition_penalty 1.1\n |> with_max_new_tokens 256\n\nlet from_preset = function\n | \"creative_writing\" -> creative_writing\n | \"chat\" -> chat\n | \"code_generation\" -> code_generation\n | \"factual\" -> factual\n | _ -> default\n\n(* ───── Logits Processors ───── *)\n\nlet neg_infinity = Float.neg_infinity\n\nlet temperature_warper ~temperature =\n {\n name = Printf.sprintf \"temperature(%.2f)\" temperature;\n process =\n (fun ~prompt_length:_ _tokens logits ->\n if temperature = 1.0 then logits else Nx.div_s logits temperature);\n }\n\nlet top_k_warper ~k =\n {\n name = Printf.sprintf \"top_k(%d)\" k;\n process =\n (fun ~prompt_length:_ _tokens logits ->\n if k <= 0 then logits\n else\n let sorted_values, _sorted_indices =\n Nx.sort ~descending:true logits\n in\n let vocab_size = Nx.numel logits in\n let cutoff_k = min k vocab_size in\n let threshold = Nx.item [ cutoff_k - 1 ] sorted_values in\n let mask = Nx.less_s logits threshold in\n Nx.where mask (Nx.full_like logits neg_infinity) logits);\n }\n\nlet top_p_warper ~p =\n {\n name = Printf.sprintf \"top_p(%.2f)\" p;\n process =\n (fun ~prompt_length:_ _tokens logits ->\n if p >= 1.0 then logits\n else\n let probs = Nx.softmax logits in\n let sorted_probs, sorted_indices = Nx.sort ~descending:true probs in\n let cumulative = Nx.cumsum sorted_probs in\n (* Find where cumulative exceeds p, keeping at least 1 token *)\n let cutoff_mask = Nx.greater_s cumulative p in\n (* Shift mask right by 1 so the token that crosses p is kept *)\n let n = Nx.numel logits in\n let shifted_arr = Nx.to_array cutoff_mask in\n let new_mask_arr = Array.make n false in\n for i = 1 to n - 1 do\n new_mask_arr.(i) <- shifted_arr.(i - 1)\n done;\n let shifted_mask = Nx.create Nx.bool [| n |] new_mask_arr in\n (* Map mask back to original token order *)\n let result = Nx.copy logits in\n let sorted_idx_arr = Nx.to_array sorted_indices in\n let shifted_mask_arr = Nx.to_array shifted_mask in\n for i = 0 to n - 1 do\n if shifted_mask_arr.(i) then\n Nx.set_item\n [ Int32.to_int sorted_idx_arr.(i) ]\n neg_infinity result\n done;\n result);\n }\n\nlet repetition_penalty ~penalty =\n {\n name = Printf.sprintf \"repetition_penalty(%.2f)\" penalty;\n process =\n (fun ~prompt_length:_ previous_tokens logits ->\n if penalty = 1.0 then logits\n else\n let result = Nx.copy logits in\n let vocab_size = Nx.numel result in\n Array.iter\n (fun token_id ->\n if token_id < vocab_size then begin\n let score = Nx.item [ token_id ] result in\n let penalized =\n if score < 0.0 then score *. penalty else score /. penalty\n in\n Nx.set_item [ token_id ] penalized result\n end)\n previous_tokens;\n result);\n }\n\nlet no_repeat_ngram ~ngram_size =\n {\n name = Printf.sprintf \"no_repeat_ngram(%d)\" ngram_size;\n process =\n (fun ~prompt_length:_ previous_tokens logits ->\n let len = Array.length previous_tokens in\n if ngram_size <= 0 || len < ngram_size - 1 then logits\n else\n let result = Nx.copy logits in\n (* Get the last (ngram_size - 1) tokens as the current prefix *)\n let prefix_start = len - (ngram_size - 1) in\n let prefix =\n Array.sub previous_tokens prefix_start (ngram_size - 1)\n in\n (* Scan history for matching prefixes *)\n for i = 0 to len - ngram_size do\n let matches = ref true in\n for j = 0 to ngram_size - 2 do\n if previous_tokens.(i + j) <> prefix.(j) then matches := false\n done;\n if !matches then begin\n let blocked_token = previous_tokens.(i + ngram_size - 1) in\n if blocked_token < Nx.numel result then\n Nx.set_item [ blocked_token ] neg_infinity result\n end\n done;\n result);\n }\n\nlet min_length ~min_length ~eos_token_ids =\n {\n name = Printf.sprintf \"min_length(%d)\" min_length;\n process =\n (fun ~prompt_length:_ tokens logits ->\n if Array.length tokens >= min_length then logits\n else\n let result = Nx.copy logits in\n let vocab_size = Nx.numel result in\n List.iter\n (fun eos_id ->\n if eos_id < vocab_size then\n Nx.set_item [ eos_id ] neg_infinity result)\n eos_token_ids;\n result);\n }\n\nlet min_new_tokens ~min_new_tokens ~eos_token_ids =\n {\n name = Printf.sprintf \"min_new_tokens(%d)\" min_new_tokens;\n process =\n (fun ~prompt_length tokens logits ->\n let new_tokens = Array.length tokens - prompt_length in\n if new_tokens >= min_new_tokens then logits\n else\n let result = Nx.copy logits in\n let vocab_size = Nx.numel result in\n List.iter\n (fun eos_id ->\n if eos_id < vocab_size then\n Nx.set_item [ eos_id ] neg_infinity result)\n eos_token_ids;\n result);\n }\n\nlet bad_words ~bad_words_ids =\n {\n name = \"bad_words\";\n process =\n (fun ~prompt_length:_ tokens logits ->\n let result = Nx.copy logits in\n let len = Array.length tokens in\n let vocab_size = Nx.numel result in\n List.iter\n (fun bad_sequence ->\n let seq_len = List.length bad_sequence in\n if seq_len > 0 && len >= seq_len - 1 then (\n let prefix_len = seq_len - 1 in\n let matches = ref true in\n let prefix = List.rev (List.tl (List.rev bad_sequence)) in\n List.iteri\n (fun i expected ->\n if tokens.(len - prefix_len + i) <> expected then\n matches := false)\n prefix;\n if !matches then begin\n let bad_token = List.nth bad_sequence (seq_len - 1) in\n if bad_token < vocab_size then\n Nx.set_item [ bad_token ] neg_infinity result\n end))\n bad_words_ids;\n result);\n }\n\nlet force_words ~force_words_ids ~iteration =\n {\n name = \"force_words\";\n process =\n (fun ~prompt_length:_ _tokens logits ->\n if iteration >= List.length force_words_ids then logits\n else\n let forced_tokens = List.nth force_words_ids iteration in\n let result = Nx.full_like logits neg_infinity in\n List.iter\n (fun token_id ->\n if token_id < Nx.numel result then\n Nx.set_item [ token_id ] (Nx.item [ token_id ] logits) result)\n forced_tokens;\n result);\n }\n\nlet custom ~name ~process = { name; process }\n\n(* ───── Stopping Criteria ───── *)\n\nlet max_length_criteria ~max_length =\n {\n name = Printf.sprintf \"max_length(%d)\" max_length;\n should_stop =\n (fun ~prompt_length:_ ~start_time:_ tokens ->\n Array.length tokens >= max_length);\n }\n\nlet max_new_tokens_criteria ~max_new_tokens =\n {\n name = Printf.sprintf \"max_new_tokens(%d)\" max_new_tokens;\n should_stop =\n (fun ~prompt_length ~start_time:_ tokens ->\n Array.length tokens - prompt_length >= max_new_tokens);\n }\n\nlet eos_token_criteria ~eos_token_ids =\n {\n name = \"eos_token\";\n should_stop =\n (fun ~prompt_length:_ ~start_time:_ tokens ->\n let len = Array.length tokens in\n if len = 0 then false else List.mem tokens.(len - 1) eos_token_ids);\n }\n\nlet max_time_criteria ~max_time =\n {\n name = Printf.sprintf \"max_time(%.1fs)\" max_time;\n should_stop =\n (fun ~prompt_length:_ ~start_time _tokens ->\n Unix.gettimeofday () -. start_time > max_time);\n }\n\nlet stop_strings_criteria ~stop_strings ~decoder =\n {\n name = \"stop_strings\";\n should_stop =\n (fun ~prompt_length:_ ~start_time:_ tokens ->\n let text = decoder tokens in\n List.exists\n (fun stop_str -> String_util.contains_substring text stop_str)\n stop_strings);\n }\n\nlet custom_criteria ~name ~should_stop = { name; should_stop }\n\n(* ───── Utilities ───── *)\n\nlet apply_processors ~processors ~prompt_length ~tokens ~logits =\n List.fold_left\n (fun acc processor -> processor.process ~prompt_length tokens acc)\n logits processors\n\nlet check_stopping ~criteria ~prompt_length ~start_time ~tokens =\n List.exists\n (fun criterion -> criterion.should_stop ~prompt_length ~start_time tokens)\n criteria\n\n(* ───── Main Generation Functions ───── *)\n\ntype generation_output = {\n sequences : int array list;\n scores : float list list option;\n}\n\nlet sample_from_logits logits =\n let probs = Nx.softmax logits in\n let probs_arr = Nx.to_array probs in\n let r = Random.float 1.0 in\n let cumsum = ref 0.0 in\n let result = ref 0 in\n (try\n for i = 0 to Array.length probs_arr - 1 do\n cumsum := !cumsum +. probs_arr.(i);\n if !cumsum > r then begin\n result := i;\n raise_notrace Exit\n end\n done\n with Exit -> ());\n !result\n\nlet argmax logits = Int32.to_int (Nx.item [ 0 ] (Nx.argmax logits))\n\nlet generate ~model ?(input_ids = [||]) ?(generation_config = default)\n ?(logits_processor = []) ?(stopping_criteria = []) () =\n let start_time = Unix.gettimeofday () in\n let prompt_length = Array.length input_ids in\n\n let processors =\n let ps = [] in\n let ps =\n if generation_config.temperature <> 1.0 then\n temperature_warper ~temperature:generation_config.temperature :: ps\n else ps\n in\n let ps =\n if generation_config.top_k > 0 then\n top_k_warper ~k:generation_config.top_k :: ps\n else ps\n in\n let ps =\n if generation_config.top_p < 1.0 then\n top_p_warper ~p:generation_config.top_p :: ps\n else ps\n in\n let ps =\n if generation_config.repetition_penalty <> 1.0 then\n repetition_penalty ~penalty:generation_config.repetition_penalty :: ps\n else ps\n in\n let ps =\n if generation_config.no_repeat_ngram_size > 0 then\n no_repeat_ngram ~ngram_size:generation_config.no_repeat_ngram_size :: ps\n else ps\n in\n let eos_ids =\n match generation_config.eos_token_id with\n | Some id -> id :: generation_config.eos_token_ids\n | None -> generation_config.eos_token_ids\n in\n let ps =\n if generation_config.min_length > 0 then\n min_length ~min_length:generation_config.min_length\n ~eos_token_ids:eos_ids\n :: ps\n else ps\n in\n let ps =\n if generation_config.min_new_tokens > 0 then\n min_new_tokens ~min_new_tokens:generation_config.min_new_tokens\n ~eos_token_ids:eos_ids\n :: ps\n else ps\n in\n ps @ logits_processor\n in\n\n let criteria =\n let cs = [] in\n let cs =\n max_length_criteria ~max_length:generation_config.max_length :: cs\n in\n let cs =\n match generation_config.max_new_tokens with\n | Some max_new -> max_new_tokens_criteria ~max_new_tokens:max_new :: cs\n | None -> cs\n in\n let eos_ids =\n match generation_config.eos_token_id with\n | Some id -> id :: generation_config.eos_token_ids\n | None -> generation_config.eos_token_ids\n in\n let cs =\n if eos_ids <> [] then eos_token_criteria ~eos_token_ids:eos_ids :: cs\n else cs\n in\n cs @ stopping_criteria\n in\n\n let tokens_ref = ref (Array.copy input_ids) in\n\n let rec generate_loop () =\n let current_tokens = !tokens_ref in\n if\n Array.length current_tokens > prompt_length\n && check_stopping ~criteria ~prompt_length ~start_time\n ~tokens:current_tokens\n then current_tokens\n else begin\n let raw_logits = model current_tokens in\n let processed =\n apply_processors ~processors ~prompt_length ~tokens:current_tokens\n ~logits:raw_logits\n in\n let next_token =\n if generation_config.do_sample then sample_from_logits processed\n else argmax processed\n in\n tokens_ref := Array.append current_tokens [| next_token |];\n generate_loop ()\n end\n in\n\n let sequences = generate_loop () in\n { sequences = [ sequences ]; scores = None }\n\nlet generate_text ~model ~tokenizer ~decoder ?(prompt = \"\")\n ?(generation_config = default) ?(logits_processor = [])\n ?(stopping_criteria = []) () =\n let input_ids = tokenizer prompt in\n let output =\n generate ~model ~input_ids ~generation_config ~logits_processor\n ~stopping_criteria ()\n in\n match output.sequences with seq :: _ -> decoder seq | [] -> \"\"\n"
},
{
"path": "dev/mimir/lib/sampler.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Text generation with composable logits processors.\n\n Provides the autoregressive decode loop, composable logits processors,\n stopping criteria, and generation configuration for language model\n inference. Operates on nx tensors for logits. *)\n\n(** {1 Core Types} *)\n\ntype logits = (float, Bigarray.float32_elt) Nx.t\n(** 1D float32 tensor of unnormalized token probabilities. Length equals\n vocabulary size. *)\n\ntype token_ids = int array\n(** Sequence of token IDs representing encoded text. *)\n\ntype logits_processor = {\n name : string;\n process : prompt_length:int -> token_ids -> logits -> logits;\n}\n(** Transforms logits before sampling. *)\n\ntype logits_processor_list = logits_processor list\n\ntype stopping_criterion = {\n name : string;\n should_stop : prompt_length:int -> start_time:float -> token_ids -> bool;\n}\n(** Determines when to end generation. *)\n\ntype stopping_criteria_list = stopping_criterion list\n\n(** {1 Generation Configuration} *)\n\ntype generation_config = {\n max_length : int;\n max_new_tokens : int option;\n min_length : int;\n min_new_tokens : int;\n do_sample : bool;\n temperature : float;\n top_k : int;\n top_p : float;\n repetition_penalty : float;\n no_repeat_ngram_size : int;\n bad_words_ids : int list list;\n force_words_ids : int list list;\n pad_token_id : int option;\n bos_token_id : int option;\n eos_token_id : int option;\n eos_token_ids : int list;\n}\n\nval default : generation_config\n\n(** {2 Builder Pattern} *)\n\nval with_temperature : float -> generation_config -> generation_config\nval with_top_k : int -> generation_config -> generation_config\nval with_top_p : float -> generation_config -> generation_config\nval with_repetition_penalty : float -> generation_config -> generation_config\nval with_max_length : int -> generation_config -> generation_config\nval with_max_new_tokens : int -> generation_config -> generation_config\nval with_min_length : int -> generation_config -> generation_config\nval with_min_new_tokens : int -> generation_config -> generation_config\nval with_no_repeat_ngram : int -> generation_config -> generation_config\nval with_do_sample : bool -> generation_config -> generation_config\n\n(** {2 Presets} *)\n\nval creative_writing : generation_config\nval chat : generation_config\nval code_generation : generation_config\nval factual : generation_config\nval from_preset : string -> generation_config\n\n(** {1 Logits Processors} *)\n\nval temperature_warper : temperature:float -> logits_processor\nval top_k_warper : k:int -> logits_processor\nval top_p_warper : p:float -> logits_processor\nval repetition_penalty : penalty:float -> logits_processor\nval no_repeat_ngram : ngram_size:int -> logits_processor\nval min_length : min_length:int -> eos_token_ids:int list -> logits_processor\n\nval min_new_tokens :\n min_new_tokens:int -> eos_token_ids:int list -> logits_processor\n\nval bad_words : bad_words_ids:int list list -> logits_processor\n\nval force_words :\n force_words_ids:int list list -> iteration:int -> logits_processor\n\nval custom :\n name:string ->\n process:(prompt_length:int -> token_ids -> logits -> logits) ->\n logits_processor\n\n(** {1 Stopping Criteria} *)\n\nval max_length_criteria : max_length:int -> stopping_criterion\nval max_new_tokens_criteria : max_new_tokens:int -> stopping_criterion\nval eos_token_criteria : eos_token_ids:int list -> stopping_criterion\nval max_time_criteria : max_time:float -> stopping_criterion\n\nval stop_strings_criteria :\n stop_strings:string list ->\n decoder:(token_ids -> string) ->\n stopping_criterion\n\nval custom_criteria :\n name:string ->\n should_stop:(prompt_length:int -> start_time:float -> token_ids -> bool) ->\n stopping_criterion\n\n(** {1 Generation} *)\n\ntype generation_output = {\n sequences : int array list;\n scores : float list list option;\n}\n\nval generate :\n model:(token_ids -> logits) ->\n ?input_ids:token_ids ->\n ?generation_config:generation_config ->\n ?logits_processor:logits_processor_list ->\n ?stopping_criteria:stopping_criteria_list ->\n unit ->\n generation_output\n\nval generate_text :\n model:(token_ids -> logits) ->\n tokenizer:(string -> token_ids) ->\n decoder:(token_ids -> string) ->\n ?prompt:string ->\n ?generation_config:generation_config ->\n ?logits_processor:logits_processor_list ->\n ?stopping_criteria:stopping_criteria_list ->\n unit ->\n string\n\n(** {1 Utilities} *)\n\nval apply_processors :\n processors:logits_processor_list ->\n prompt_length:int ->\n tokens:token_ids ->\n logits:logits ->\n logits\n\nval check_stopping :\n criteria:stopping_criteria_list ->\n prompt_length:int ->\n start_time:float ->\n tokens:token_ids ->\n bool\n"
},
{
"path": "dev/mimir/lib/string_util.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet contains_substring s sub =\n let len_s = String.length s in\n let len_sub = String.length sub in\n if len_sub = 0 then true\n else if len_sub > len_s then false\n else\n let rec check i =\n if i > len_s - len_sub then false\n else if String.sub s i len_sub = sub then true\n else check (i + 1)\n in\n check 0\n"
},
{
"path": "dev/umbra/README.md",
"content": "# Umbra\n\nComputational astronomy for OCaml, powered by [Nx](../../packages/nx/) and [Rune](../../packages/rune/)\n\nUmbra provides dimensionally-typed physical quantities, cosmological distances,\nspectral energy distributions, dust extinction, synthetic photometry, coordinate\ntransforms, time scales, catalog cross-matching, and weak lensing survey science.\nAll computations operate on Nx tensors and are differentiable through Rune --\nfit cosmological parameters, propagate uncertainties via Jacobians, or sample\nposteriors with HMC, all from the same forward model.\n\n## Quick Start\n\nCompute the luminosity distance to a galaxy at redshift 0.5:\n\n```ocaml\nopen Umbra\n\nlet () =\n let f64 = Nx.float64 in\n let z = Nx.scalar f64 0.5 in\n let dl = Cosmo.luminosity_distance ~p:Cosmo.planck18 z in\n Printf.printf \"d_L(z=0.5) = %.1f Mpc\\n\"\n (Nx.item [] (Unit.Length.in_mpc dl))\n```\n\nFit stellar temperature from photometry with automatic derivatives:\n\n```ocaml\nlet model params =\n let temp = Unit.Temperature.of_kelvin (Nx.exp (Nx.slice [ I 0 ] params)) in\n let av = Nx.reshape [||] (Nx.slice [ I 1 ] params) in\n let rv = Nx.scalar Nx.float64 3.1 in\n List.map (fun bp ->\n let wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:wave\n |> Extinction.apply (Extinction.ccm89 ~rv) ~av\n |> Spectrum.as_flux_density\n in\n Photometry.ab_mag bp sed) bands\n |> Nx.stack ~axis:0\n\n(* Rune differentiates through the entire pipeline *)\nlet loss, grad = Rune.value_and_grad chi2 params\n```\n\n## Features\n\n- **Dimensional types**: `Unit.Length`, `Unit.Mass`, `Unit.Time`, `Unit.Angle`, etc. with compile-time safety\n- **Physical constants**: CODATA 2022 and IAU 2015 via `Const`\n- **Cosmology**: LCDM, wCDM, w0waCDM distances, growth factors, and matter power spectra via `Cosmo`\n- **Spectra**: blackbody, power-law, and line profiles (Gaussian, Lorentzian, Voigt) via `Spectrum`\n- **Extinction**: CCM89, Fitzpatrick99, O'Donnell94, Calzetti00 dust laws via `Extinction`\n- **Photometry**: AB, ST, and Vega magnitudes through standard filter bandpasses via `Photometry`\n- **Filters**: SDSS, Johnson-Cousins, 2MASS, Gaia DR3, Rubin/LSST, Euclid via `Filters`\n- **Coordinates**: ICRS, Galactic, Ecliptic, Supergalactic frame transforms and kd-tree cross-matching via `Coord`\n- **Time**: UTC, TAI, TT, TDB time scales with phantom-typed safety via `Time`\n- **Observer geometry**: altitude-azimuth coordinates and airmass via `Altaz`\n- **Survey science**: angular power spectra and Fisher forecasting via `Survey`\n- **FITS I/O**: image and table read/write via `Umbra_fits`\n- **Fully differentiable**: all forward models work with Rune's autodiff, Jacobians, and MCMC\n\n## Examples\n\n| Example | Concept |\n|---------|---------|\n| [`01-constants-and-units`](examples/01-constants-and-units/) | Type-safe physical quantities and conversions |\n| [`02-cosmological-distances`](examples/02-cosmological-distances/) | LCDM distances and SN Ia fitting |\n| [`03-blackbody-fitting`](examples/03-blackbody-fitting/) | Fit stellar temperature from photometry |\n| [`04-extinction-and-magnitudes`](examples/04-extinction-and-magnitudes/) | Dust extinction, magnitude systems, K-corrections |\n| [`05-sed-fitting`](examples/05-sed-fitting/) | Full SED pipeline: blackbody, extinction, photometry |\n| [`06-coordinates-and-time`](examples/06-coordinates-and-time/) | Frame transforms, time scales, observer geometry |\n| [`07-batch-photometry`](examples/07-batch-photometry/) | Batched operations over parameter grids |\n| [`08-photometric-redshifts`](examples/08-photometric-redshifts/) | Two-stage photo-z: grid search + gradient refinement |\n| [`09-gravitational-lensing`](examples/09-gravitational-lensing/) | Point-mass lens model parameter fitting |\n| [`10-uncertainty-propagation`](examples/10-uncertainty-propagation/) | AD Jacobians for error propagation vs Monte Carlo |\n| [`11-bayesian-sed`](examples/11-bayesian-sed/) | Fisher matrix + HMC posterior sampling |\n| [`12-survey-optimization`](examples/12-survey-optimization/) | Differentiable Fisher forecasting for survey design |\n\n## Papers\n\n- [**Perlmutter et al. 1999**](papers/perlmutter1999/) -- Reproducing the Nobel Prize-winning discovery of cosmic acceleration using the Pantheon+ dataset\n\n## Contributing\n\nSee the [Raven monorepo README](../../README.md) for guidelines.\n\n## License\n\nISC License. See [LICENSE](../../LICENSE) for details.\n"
},
{
"path": "dev/umbra/dune-project",
"content": "(lang dune 3.21)\n\n(name umbra)\n\n(package\n (name umbra)\n (synopsis \"Astronomy library for OCaml\")\n (description\n \"Physical units, celestial coordinates, FITS I/O, cosmological distances, and catalog cross-matching. Built on Nx and Talon.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (nx\n (>= 1.0.0~alpha3))\n (talon\n (>= 1.0.0~alpha3))\n (windtrap :with-test)))\n"
},
{
"path": "dev/umbra/examples/01-constants-and-units/README.md",
"content": "# `01-constants-and-units`\n\nIntroduction to Umbra's type-safe unit system and physical constants. Creates\nquantities in different units, converts between them, and demonstrates how\nphantom types prevent mixing incompatible dimensions at compile time.\n\n```bash\ndune exec dev/umbra/examples/01-constants-and-units/main.exe\n```\n\n## What You'll Learn\n\n- Creating quantities with scalar constructors (`Length.pc`, `Angle.deg`, `Mass.solar_mass`)\n- Converting between units (`Length.in_ly`, `Angle.in_arcsec`)\n- Adding quantities of the same dimension (`Unit.(+)`)\n- Using physical constants (`Const.c`, `Const.h_si`, `Const.k_b_si`)\n- Cross-dimension conversions (`parallax_to_distance`, `wavelength_to_frequency`)\n- Batch operations on tensor-valued quantities\n\n## Key Functions\n\n| Function | Purpose |\n| --------------------------- | -------------------------------------------- |\n| `Length.pc`, `Length.au` | Create length quantities in parsecs, AU |\n| `Length.in_m`, `Length.in_ly`| Extract values in metres, light-years |\n| `Angle.deg`, `Angle.arcsec` | Create angles in degrees, arcseconds |\n| `Temperature.kelvin` | Create temperature quantities |\n| `Mass.solar_mass` | Create mass in solar masses |\n| `Const.c`, `Const.h_si` | Speed of light, Planck constant |\n| `parallax_to_distance` | Convert stellar parallax to distance |\n| `wavelength_to_frequency` | Convert wavelength to frequency via c/lambda |\n\n## Try It\n\n1. Compute the Schwarzschild radius of the Sun using `Const.g_si`, `Const.solar_mass`, and `Const.c`.\n2. Add `Length.ly 4.246` (Proxima Centauri) and check it matches the parallax-derived distance.\n3. Use `Unit.doppler_optical` to compute the observed wavelength of H-alpha at a radial velocity of 100 km/s.\n\n## Next Steps\n\nContinue to [02-cosmological-distances](../02-cosmological-distances/) to compute\ndistances and times in an expanding universe.\n"
},
{
"path": "dev/umbra/examples/01-constants-and-units/dune",
"content": "(executable\n (name main)\n (libraries nx umbra))\n"
},
{
"path": "dev/umbra/examples/01-constants-and-units/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Type-safe units and physical constants.\n\n Introduces Umbra's dimensional type system: quantities carry phantom types\n that prevent mixing incompatible dimensions at compile time. Shows how to\n create, convert, and combine quantities in different units, and how to use\n physical and astronomical constants. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\nlet () =\n Printf.printf \"Type-safe units and physical constants\\n\";\n Printf.printf \"======================================\\n\\n\";\n\n (* --- Length: metres, parsecs, AU, light-years --- *)\n Printf.printf \"Length conversions\\n\";\n Printf.printf \"------------------\\n\";\n let d_pc = Unit.Length.pc 1.0 in\n Printf.printf \" 1 parsec = %.4e m\\n\" (item [] (Unit.Length.in_m d_pc));\n Printf.printf \" 1 parsec = %.6f ly\\n\" (item [] (Unit.Length.in_ly d_pc));\n Printf.printf \" 1 parsec = %.0f AU\\n\" (item [] (Unit.Length.in_au d_pc));\n\n let d_au = Unit.Length.au 1.0 in\n Printf.printf \" 1 AU = %.4e m\\n\" (item [] (Unit.Length.in_m d_au));\n Printf.printf \" 1 AU = %.4e pc\\n\\n\" (item [] (Unit.Length.in_pc d_au));\n\n (* Adding lengths of different units — the type system ensures consistency *)\n let d_total = Unit.( + ) (Unit.Length.kpc 10.0) (Unit.Length.pc 500.0) in\n Printf.printf \" 10 kpc + 500 pc = %.3f kpc\\n\\n\"\n (item [] (Unit.Length.in_kpc d_total));\n\n (* --- Angle: degrees, radians, arcseconds --- *)\n Printf.printf \"Angle conversions\\n\";\n Printf.printf \"-----------------\\n\";\n let a_deg = Unit.Angle.deg 1.0 in\n Printf.printf \" 1 degree = %.6f rad\\n\" (item [] (Unit.Angle.in_rad a_deg));\n Printf.printf \" 1 degree = %.1f arcmin\\n\"\n (item [] (Unit.Angle.in_arcmin a_deg));\n Printf.printf \" 1 degree = %.1f arcsec\\n\"\n (item [] (Unit.Angle.in_arcsec a_deg));\n\n let a_mas = Unit.Angle.mas 1.0 in\n Printf.printf \" 1 mas = %.4e arcsec\\n\\n\"\n (item [] (Unit.Angle.in_arcsec a_mas));\n\n (* --- Temperature --- *)\n Printf.printf \"Temperature\\n\";\n Printf.printf \"-----------\\n\";\n let sun_t = Unit.Temperature.kelvin 5778.0 in\n Printf.printf \" Sun surface: %.0f K\\n\"\n (item [] (Unit.Temperature.in_kelvin sun_t));\n let sirius_t = Unit.Temperature.kelvin 9940.0 in\n Printf.printf \" Sirius: %.0f K\\n\\n\"\n (item [] (Unit.Temperature.in_kelvin sirius_t));\n\n (* --- Time durations --- *)\n Printf.printf \"Time durations\\n\";\n Printf.printf \"--------------\\n\";\n let t_yr = Unit.Time.yr 1.0 in\n Printf.printf \" 1 Julian year = %.0f days\\n\"\n (item [] (Unit.Time.in_day t_yr));\n Printf.printf \" 1 Julian year = %.2f s\\n\" (item [] (Unit.Time.in_s t_yr));\n\n let t_gyr = Unit.Time.gyr 13.8 in\n Printf.printf \" Age of universe ~ %.2e yr\\n\\n\"\n (item [] (Unit.Time.in_yr t_gyr));\n\n (* --- Mass: kg, solar masses, Earth masses --- *)\n Printf.printf \"Mass conversions\\n\";\n Printf.printf \"----------------\\n\";\n let m_sun = Unit.Mass.solar_mass 1.0 in\n Printf.printf \" 1 solar mass = %.4e kg\\n\" (item [] (Unit.Mass.in_kg m_sun));\n Printf.printf \" 1 solar mass = %.0f Earth masses\\n\"\n (item [] (Unit.Mass.in_earth_mass m_sun));\n Printf.printf \" 1 solar mass = %.1f Jupiter masses\\n\\n\"\n (item [] (Unit.Mass.in_jupiter_mass m_sun));\n\n (* --- Physical constants --- *)\n Printf.printf \"Physical constants\\n\";\n Printf.printf \"------------------\\n\";\n Printf.printf \" c = %.0f m/s\\n\" (Unit.to_float Const.c);\n Printf.printf \" h = %.4e J s\\n\" Const.h_si;\n Printf.printf \" k_B = %.4e J/K\\n\" Const.k_b_si;\n Printf.printf \" G = %.4e m^3 kg^-1 s^-2\\n\" Const.g_si;\n Printf.printf \" sigma = %.4e W m^-2 K^-4\\n\\n\" Const.sigma_sb_si;\n\n (* --- Astronomical constants --- *)\n Printf.printf \"Astronomical constants\\n\";\n Printf.printf \"----------------------\\n\";\n Printf.printf \" L_sun = %.4e W\\n\"\n (item [] (Unit.Power.in_w Const.solar_luminosity));\n Printf.printf \" R_sun = %.4e m\\n\"\n (item [] (Unit.Length.in_m Const.solar_radius));\n Printf.printf \" M_sun = %.4e kg\\n\"\n (item [] (Unit.Mass.in_kg Const.solar_mass));\n Printf.printf \" 1 AU = %.4e m\\n\" (item [] (Unit.Length.in_m Const.au));\n Printf.printf \" 1 pc = %.4e m\\n\\n\" (item [] (Unit.Length.in_m Const.pc));\n\n (* --- Cross-dimension: parallax to distance --- *)\n Printf.printf \"Parallax to distance\\n\";\n Printf.printf \"--------------------\\n\";\n let parallax = Unit.Angle.arcsec 1.0 in\n let dist = Unit.parallax_to_distance parallax in\n Printf.printf \" 1 arcsec parallax -> %.3f pc\\n\"\n (item [] (Unit.Length.in_pc dist));\n\n let proxima_parallax = Unit.Angle.mas 768.5 in\n let proxima_dist = Unit.parallax_to_distance proxima_parallax in\n Printf.printf \" Proxima Cen (768.5 mas) -> %.3f pc\\n\"\n (item [] (Unit.Length.in_pc proxima_dist));\n\n (* --- Tensor operations: batch unit conversions --- *)\n Printf.printf \"\\nBatch operations\\n\";\n Printf.printf \"----------------\\n\";\n let wavelengths_nm =\n create f64 [| 5 |] [| 380.0; 450.0; 550.0; 650.0; 750.0 |]\n in\n let wavelengths = Unit.Length.of_nm wavelengths_nm in\n let wavelengths_angstrom = Unit.Length.in_angstrom wavelengths in\n Printf.printf \" Wavelengths (nm): %s\\n\"\n (Nx.data_to_string wavelengths_nm);\n Printf.printf \" Wavelengths (angstrom): %s\\n\"\n (Nx.data_to_string wavelengths_angstrom);\n\n (* Convert wavelength to frequency *)\n let freqs = Unit.wavelength_to_frequency wavelengths in\n Printf.printf \" Frequencies (Hz): %s\\n\"\n (Nx.data_to_string (Unit.Frequency.in_hz freqs));\n\n (* Wien's law: peak wavelength of a blackbody *)\n Printf.printf \"\\nWien's displacement law\\n\";\n Printf.printf \"----------------------\\n\";\n let b_wien = Const.b_wien_si in\n let sun_peak_m = b_wien /. item [] (Unit.Temperature.in_kelvin sun_t) in\n Printf.printf \" Sun (T=%.0f K): peak at %.0f nm\\n\"\n (item [] (Unit.Temperature.in_kelvin sun_t))\n (sun_peak_m *. 1e9);\n let sirius_peak_m = b_wien /. item [] (Unit.Temperature.in_kelvin sirius_t) in\n Printf.printf \" Sirius (T=%.0f K): peak at %.0f nm\\n\"\n (item [] (Unit.Temperature.in_kelvin sirius_t))\n (sirius_peak_m *. 1e9)\n"
},
{
"path": "dev/umbra/examples/02-cosmological-distances/README.md",
"content": "# `02-cosmological-distances`\n\nCosmological distance calculations and parameter fitting. First prints a\ndistance table for the Planck 2018 cosmology, then fits H0 and Omega_m from\nsynthetic Type Ia supernova distance moduli using gradient descent.\n\n```bash\ndune exec dev/umbra/examples/02-cosmological-distances/main.exe\n```\n\n## What You'll Learn\n\n- Using preset cosmologies (`Cosmo.planck18`)\n- Computing distances (`comoving_distance`, `luminosity_distance`, `angular_diameter_distance`)\n- Computing distance modulus and lookback time\n- Building differentiable cosmological models with `create_flat_lcdm`\n- Fitting cosmological parameters with Rune autodiff and Vega optimizers\n\n## Key Functions\n\n| Function | Purpose |\n| ----------------------------- | --------------------------------------------- |\n| `Cosmo.planck18` | Planck 2018 flat LCDM preset |\n| `Cosmo.comoving_distance` | Line-of-sight comoving distance |\n| `Cosmo.luminosity_distance` | Luminosity distance at redshift z |\n| `Cosmo.distance_modulus` | Distance modulus mu = 5 log10(d_L/Mpc) + 25 |\n| `Cosmo.lookback_time` | Time since light was emitted |\n| `Cosmo.age` | Age of the universe at redshift z |\n| `Cosmo.create_flat_lcdm` | Tensor-parameterized cosmology for autodiff |\n| `Rune.value_and_grads` | Forward pass + gradient computation |\n\n## How It Works\n\nThe distance modulus forward model uses `Cosmo.distance_modulus`, which\ninternally integrates E(z) via 16-point Gauss-Legendre quadrature. Since all\noperations are Nx tensor ops, gradients flow through the entire pipeline\nautomatically via Rune.\n\nThe optimizer starts from H0=65, Omega_m=0.25 and converges toward the true\nvalues (H0~73, Omega_m~0.3) that generated the synthetic data.\n\n## Try It\n\n1. Change the preset to `Cosmo.wmap9` and compare the distance table.\n2. Add `Omega_L` as a free parameter using `create_lcdm` for a non-flat model.\n3. Use `Cosmo.z_at_value` to find the redshift where the lookback time is 10 Gyr.\n\n## Next Steps\n\nContinue to [03-blackbody-fitting](../03-blackbody-fitting/) to fit stellar\ntemperatures from photometry.\n"
},
{
"path": "dev/umbra/examples/02-cosmological-distances/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega umbra))\n"
},
{
"path": "dev/umbra/examples/02-cosmological-distances/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Differentiable cosmological parameter fitting from Type Ia supernova distance\n moduli.\n\n Fits H0 (Hubble constant) and Omega_m (matter density fraction) by gradient\n descent on the distance modulus residuals. The forward model uses\n Umbra.Cosmo.distance_modulus directly -- its Gauss-Legendre quadrature,\n luminosity distance, and distance modulus are all Nx tensor operations,\n making them natively differentiable through Rune's autodiff.\n\n Also demonstrates basic cosmological distance queries: comoving distance,\n luminosity distance, angular diameter distance, lookback time, and the age of\n the universe at various redshifts. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\n(* --- Part 1: Distance table for the Planck 2018 cosmology --- *)\n\nlet print_distance_table () =\n Printf.printf \"Cosmological distances (Planck 2018)\\n\";\n Printf.printf \"====================================\\n\\n\";\n let p = Cosmo.planck18 in\n Printf.printf \" H0 = %.2f km/s/Mpc\\n\" (item [] (Cosmo.h0 p));\n Printf.printf \" Omega_m = %.4f\\n\" (item [] (Cosmo.omega_m p));\n Printf.printf \" Omega_L = %.4f\\n\\n\" (item [] (Cosmo.omega_l p));\n\n Printf.printf \"%6s %10s %10s %10s %8s %8s\\n\" \"z\" \"d_C (Mpc)\" \"d_L (Mpc)\"\n \"d_A (Mpc)\" \"mu\" \"t_lb (Gyr)\";\n Printf.printf \"%6s %10s %10s %10s %8s %8s\\n\" \"------\" \"----------\"\n \"----------\" \"----------\" \"--------\" \"----------\";\n let redshifts = [| 0.01; 0.1; 0.3; 0.5; 1.0; 2.0; 3.0; 5.0 |] in\n Array.iter\n (fun z ->\n let zv = scalar f64 z in\n let d_c = item [] (Unit.Length.in_mpc (Cosmo.comoving_distance ~p zv)) in\n let d_l =\n item [] (Unit.Length.in_mpc (Cosmo.luminosity_distance ~p zv))\n in\n let d_a =\n item [] (Unit.Length.in_mpc (Cosmo.angular_diameter_distance ~p zv))\n in\n let mu = item [] (Cosmo.distance_modulus ~p zv) in\n let t_lb = item [] (Unit.Time.in_gyr (Cosmo.lookback_time ~p zv)) in\n Printf.printf \"%6.2f %10.1f %10.1f %10.1f %8.2f %8.2f\\n\" z d_c d_l\n d_a mu t_lb)\n redshifts;\n Printf.printf \"\\n\";\n\n (* Age of the universe *)\n let age_now = item [] (Unit.Time.in_gyr (Cosmo.age ~p (scalar f64 0.0))) in\n Printf.printf \" Age of the universe (z=0): %.2f Gyr\\n\\n\" age_now\n\n(* --- Part 2: Fit H0 and Omega_m from SN Ia data --- *)\n\n(* Representative SN Ia data points (z, observed distance modulus). Based on\n Pantheon+ compilation values for flat LCDM with H0 ~ 73, Omega_m ~ 0.3. *)\nlet z_arr = [| 0.01; 0.03; 0.08; 0.15; 0.25; 0.40; 0.55; 0.70; 0.85; 1.00 |]\nlet n_sn = Array.length z_arr\n\nlet mu_obs =\n [| 33.07; 35.47; 37.62; 39.07; 40.24; 41.42; 42.23; 42.85; 43.34; 43.74 |]\n\n(* Forward model: compute distance modulus for all SNe. The differentiable\n parameters are H0 and Omega_m, which flow through Cosmo.distance_modulus via\n Nx tensor operations. *)\nlet loss params =\n match params with\n | [ h0; omega_m ] ->\n let p = Cosmo.create_flat_lcdm ~h0 ~omega_m in\n let total = ref (scalar f64 0.0) in\n for i = 0 to n_sn - 1 do\n let z_i = scalar f64 z_arr.(i) in\n let mu_pred = Cosmo.distance_modulus ~p z_i in\n let mu_obs_i = scalar f64 mu_obs.(i) in\n let residual = sub mu_pred mu_obs_i in\n total := add !total (square residual)\n done;\n div_s !total (Float.of_int n_sn)\n | _ -> failwith \"expected [h0; omega_m]\"\n\nlet fit_cosmology () =\n Printf.printf \"Fitting H0 and Omega_m from Type Ia supernovae\\n\";\n Printf.printf \"===============================================\\n\";\n Printf.printf \" Data: %d distance moduli (Pantheon+-like)\\n\" n_sn;\n Printf.printf \" Method: Adam optimizer, 300 steps\\n\";\n Printf.printf \" Model: flat LCDM via Cosmo.distance_modulus\\n\\n\";\n\n let algo = Vega.adam (Vega.Schedule.constant 0.5) in\n let h0 = ref (scalar f64 65.0) in\n let omega_m = ref (scalar f64 0.25) in\n let states = [| Vega.init algo !h0; Vega.init algo !omega_m |] in\n let steps = 300 in\n\n Printf.printf \"%5s %10s %8s %8s\\n\" \"step\" \"loss\" \"H0\" \"Omega_m\";\n Printf.printf \"%5s %10s %8s %8s\\n\" \"-----\" \"----------\" \"--------\"\n \"--------\";\n\n let refs = [| h0; omega_m |] in\n for i = 0 to steps - 1 do\n let loss_val, grads = Rune.value_and_grads loss [ !h0; !omega_m ] in\n List.iteri\n (fun j g ->\n let p, s = Vega.step states.(j) ~grad:g ~param:!(refs.(j)) in\n refs.(j) := p;\n states.(j) <- s)\n grads;\n if i mod 50 = 0 || i = steps - 1 then\n Printf.printf \"%5d %10.6f %8.2f %8.4f\\n\" i (item [] loss_val)\n (item [] !h0) (item [] !omega_m)\n done;\n\n Printf.printf \"\\nFitted parameters:\\n\";\n Printf.printf \" H0 = %.2f km/s/Mpc\\n\" (item [] !h0);\n Printf.printf \" Omega_m = %.4f\\n\" (item [] !omega_m)\n\nlet () =\n print_distance_table ();\n fit_cosmology ()\n"
},
{
"path": "dev/umbra/examples/03-blackbody-fitting/README.md",
"content": "# `03-blackbody-fitting`\n\nFits the effective temperature and luminosity normalization of a star from\nsynthetic UGRIZ broadband photometry using gradient descent on a blackbody\nmodel.\n\n```bash\ndune exec dev/umbra/examples/03-blackbody-fitting/main.exe\n```\n\n## What You'll Learn\n\n- Using physical constants (`Const.h_si`, `Const.k_b_si`, `Const.c`)\n- Building a differentiable Planck function from Nx tensor operations\n- Parameterizing in log-space for numerical stability\n- Fitting chi-squared with Rune autodiff and Vega's Adam optimizer\n\n## Key Functions\n\n| Function | Purpose |\n| --------------------- | -------------------------------------------------- |\n| `Const.h_si` | Planck constant (J s) |\n| `Const.k_b_si` | Boltzmann constant (J/K) |\n| `Const.c` | Speed of light (typed velocity) |\n| `Unit.to_float` | Extract scalar SI value from a typed constant |\n| `Rune.value_and_grads`| Compute loss and gradients in one pass |\n| `Vega.adam` | Adam optimizer |\n| `Vega.step` | Apply one optimization step |\n\n## How It Works\n\nThe Planck spectral radiance B(lambda, T) = 2hc^2 / lambda^5 / (exp(hc /\nlambda k T) - 1) is implemented entirely with Nx tensor operations. Since Rune\ncan differentiate any Nx computation, gradients of chi-squared with respect to\nlog(T) and log(A) are computed automatically.\n\nThe optimizer starts from T=5000 K and converges toward the true temperature\nof 5800 K (Sun-like star). Log-space parameterization ensures positivity and\nimproves gradient conditioning.\n\n## Try It\n\n1. Change the true temperature to 10000 K (A-type star) and observe how the\n SED shape changes.\n2. Add a third parameter for a dust extinction term.\n3. Replace the central-wavelength approximation with proper filter integration\n using `Photometry.ab_mag` (see example 05).\n\n## Next Steps\n\nContinue to [04-extinction-and-magnitudes](../04-extinction-and-magnitudes/) to\nlearn about dust extinction, K-corrections, and magnitude systems.\n"
},
{
"path": "dev/umbra/examples/03-blackbody-fitting/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega umbra))\n"
},
{
"path": "dev/umbra/examples/03-blackbody-fitting/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Differentiable blackbody SED fitting.\n\n Given broadband photometric measurements in UGRIZ bands, fit the stellar\n effective temperature and luminosity normalization by gradient descent on the\n chi-squared statistic. The Planck function is evaluated as Nx tensor\n operations, making it fully differentiable through Rune.\n\n Uses Umbra.Const for physical constants. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\n(* Central wavelengths of SDSS UGRIZ bands in meters *)\nlet lambda =\n create f64 [| 5 |] [| 3.551e-7; 4.686e-7; 6.166e-7; 7.480e-7; 8.932e-7 |]\n\n(* Physical constants from Umbra *)\nlet h_planck = Const.h_si\nlet c_light = Unit.to_float Const.c\nlet k_boltz = Const.k_b_si\n\n(* Pre-computed constant tensors *)\nlet two_hc2 = scalar f64 (2.0 *. h_planck *. c_light *. c_light)\nlet hc_over_k = scalar f64 (h_planck *. c_light /. k_boltz)\nlet lam5 = pow_s lambda 5.0\n\n(* Generate synthetic observations from a Sun-like star *)\nlet true_temp = 5800.0\nlet true_log_norm = -50.0\n\nlet planck_scalar lam_m temp =\n let x = h_planck *. c_light /. (lam_m *. k_boltz *. temp) in\n 2.0 *. h_planck *. c_light *. c_light\n /. (lam_m *. lam_m *. lam_m *. lam_m *. lam_m)\n /. (Float.exp x -. 1.0)\n\nlet flux_obs =\n let norm = Float.exp true_log_norm in\n let fluxes =\n Array.init 5 (fun i ->\n let lam_m =\n [| 3.551e-7; 4.686e-7; 6.166e-7; 7.480e-7; 8.932e-7 |].(i)\n in\n norm\n *. planck_scalar lam_m true_temp\n *. (1.0 +. (0.02 *. (Float.of_int i -. 2.0))))\n in\n create f64 [| 5 |] fluxes\n\n(* Fractional errors: 5% photometry *)\nlet flux_err = mul_s flux_obs 0.05\nlet band_names = [| \"U\"; \"G\"; \"R\"; \"I\"; \"Z\" |]\n\n(* Differentiable forward model: Planck function at 5 wavelengths. Parameterized\n in log-space for positivity and gradient conditioning.\n\n B(lambda, T) = 2hc^2 / lambda^5 / (exp(hc / (lambda * k * T)) - 1) *)\nlet loss params =\n match params with\n | [ log_temp; log_norm ] ->\n let temp = exp log_temp in\n let norm = exp log_norm in\n let exponent = div hc_over_k (mul lambda temp) in\n let planck =\n div (div two_hc2 lam5) (sub (exp exponent) (scalar f64 1.0))\n in\n let flux_pred = mul norm planck in\n let residual = div (sub flux_pred flux_obs) flux_err in\n sum (square residual)\n | _ -> failwith \"expected [log_temp; log_norm]\"\n\nlet () =\n Printf.printf \"Differentiable blackbody SED fitting\\n\";\n Printf.printf \"====================================\\n\";\n Printf.printf \"Fitting temperature and normalization to UGRIZ photometry\\n\\n\";\n\n Printf.printf \"True parameters:\\n\";\n Printf.printf \" T = %.0f K\\n\" true_temp;\n Printf.printf \" logA = %.1f\\n\\n\" true_log_norm;\n\n Printf.printf \"Synthetic observations (5%% errors):\\n\";\n for i = 0 to 4 do\n Printf.printf \" %s: %.4e +/- %.4e\\n\" band_names.(i) (item [ i ] flux_obs)\n (item [ i ] flux_err)\n done;\n Printf.printf \"\\n\";\n\n (* Start from a guess *)\n let algo = Vega.adam (Vega.Schedule.constant 1e-2) in\n let log_temp = ref (scalar f64 (Float.log 5000.0)) in\n let log_norm = ref (scalar f64 (-52.0)) in\n let states = [| Vega.init algo !log_temp; Vega.init algo !log_norm |] in\n let steps = 500 in\n\n Printf.printf \"%5s %12s %8s %10s\\n\" \"step\" \"chi2\" \"T (K)\" \"log_norm\";\n Printf.printf \"%5s %12s %8s %10s\\n\" \"-----\" \"------------\" \"--------\"\n \"----------\";\n\n let refs = [| log_temp; log_norm |] in\n for i = 0 to steps - 1 do\n let loss_val, grads = Rune.value_and_grads loss [ !log_temp; !log_norm ] in\n List.iteri\n (fun j g ->\n let p, s = Vega.step states.(j) ~grad:g ~param:!(refs.(j)) in\n refs.(j) := p;\n states.(j) <- s)\n grads;\n if i mod 100 = 0 || i = steps - 1 then\n Printf.printf \"%5d %12.4f %8.1f %10.3f\\n\" i (item [] loss_val)\n (Float.exp (item [] !log_temp))\n (item [] !log_norm)\n done;\n\n Printf.printf \"\\nFitted parameters:\\n\";\n Printf.printf \" T = %.1f K (true: %.0f K)\\n\"\n (Float.exp (item [] !log_temp))\n true_temp;\n Printf.printf \" logA = %.3f (true: %.1f)\\n\" (item [] !log_norm)\n true_log_norm\n"
},
{
"path": "dev/umbra/examples/04-extinction-and-magnitudes/README.md",
"content": "# `04-extinction-and-magnitudes`\n\nExplores three key photometric concepts: magnitude systems (AB, ST, Vega),\nK-corrections from redshift, and interstellar dust extinction. Shows how to\ncompose `Spectrum`, `Extinction`, `Photometry`, and `Filters` modules.\n\n```bash\ndune exec dev/umbra/examples/04-extinction-and-magnitudes/main.exe\n```\n\n## What You'll Learn\n\n- Computing AB, ST, and Vega magnitudes through real SDSS filters\n- Understanding K-corrections from redshift-shifted SEDs\n- Applying extinction laws (CCM89, Fitzpatrick99, O'Donnell94)\n- Measuring colors and color excess from dust reddening\n\n## Key Functions\n\n| Function | Purpose |\n| --------------------------- | ---------------------------------------------- |\n| `Photometry.ab_mag` | AB magnitude through a bandpass |\n| `Photometry.st_mag` | ST magnitude through a bandpass |\n| `Photometry.vega_mag` | Vega magnitude through a bandpass |\n| `Photometry.color` | Color index (mag difference between two bands) |\n| `Spectrum.blackbody` | Planck spectral radiance |\n| `Spectrum.redshift` | Apply cosmological redshift to an SED |\n| `Spectrum.as_flux_density` | Cast spectrum to flux density kind |\n| `Extinction.ccm89` | Cardelli, Clayton & Mathis (1989) dust law |\n| `Extinction.fitzpatrick99` | Fitzpatrick (1999) dust law |\n| `Extinction.apply` | Redden a spectrum by A_V magnitudes |\n| `Filters.sdss_r` | Pre-built SDSS r-band bandpass |\n\n## How It Works\n\n**Magnitude systems** differ in their reference flux:\n- AB: constant f_nu = 3631 Jy\n- ST: constant f_lambda = 3.63e-9 erg/s/cm^2/A\n- Vega: the spectrum of alpha Lyrae\n\n**K-corrections** arise because redshift moves the SED across the bandpass,\nchanging the measured flux even without distance dimming. K(z) = m_obs - m_rest.\n\n**Extinction** attenuates and reddens starlight. The extinction curve\nA_lambda/A_V depends on wavelength and the dust grain properties (encoded in\nR_V). Higher A_V means more dimming; bluer bands are affected more, producing\nreddening.\n\n## Try It\n\n1. Compare Galactic extinction (CCM89, R_V=3.1) with starburst attenuation\n (`Extinction.calzetti00`).\n2. Apply both redshift and extinction to see their combined effect on colors.\n3. Use `Extinction.unredden` to recover the intrinsic SED from a reddened\n observation.\n\n## Next Steps\n\nContinue to [05-sed-fitting](../05-sed-fitting/) to fit temperature, extinction,\nand normalization simultaneously.\n"
},
{
"path": "dev/umbra/examples/04-extinction-and-magnitudes/dune",
"content": "(executable\n (name main)\n (libraries nx rune umbra))\n"
},
{
"path": "dev/umbra/examples/04-extinction-and-magnitudes/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* K-corrections, extinction, and magnitude systems.\n\n Demonstrates three key photometric concepts:\n\n 1. Magnitude systems: AB, ST, and Vega magnitudes through real SDSS filters.\n 2. K-corrections: the difference between observed and rest-frame magnitudes\n due to redshift shifting the SED across the bandpass. 3. Extinction: how\n interstellar dust reddens and dims stellar light, comparing CCM89 and\n Fitzpatrick99 extinction laws. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\nlet () =\n Printf.printf \"Extinction, K-corrections, and magnitude systems\\n\";\n Printf.printf \"=================================================\\n\\n\";\n\n (* --- Part 1: Magnitude systems --- *)\n Printf.printf \"Part 1: AB, ST, and Vega magnitudes\\n\";\n Printf.printf \"------------------------------------\\n\\n\";\n\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 6000.0) in\n let norm = Nx.scalar f64 (Float.exp (-50.0)) in\n let bands =\n [|\n (\"SDSS u\", Filters.sdss_u);\n (\"SDSS g\", Filters.sdss_g);\n (\"SDSS r\", Filters.sdss_r);\n (\"SDSS i\", Filters.sdss_i);\n (\"SDSS z\", Filters.sdss_z);\n |]\n in\n\n Printf.printf \" Source: T=6000 K blackbody\\n\\n\";\n Printf.printf \"%8s %8s %8s %8s\\n\" \"Band\" \"AB\" \"ST\" \"Vega\";\n Printf.printf \"%8s %8s %8s %8s\\n\" \"--------\" \"--------\" \"--------\"\n \"--------\";\n Array.iter\n (fun (name, bp) ->\n let bp_wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm |> Spectrum.as_flux_density\n in\n let m_ab = item [] (Photometry.ab_mag bp sed) in\n let m_st = item [] (Photometry.st_mag bp sed) in\n let m_vega = item [] (Photometry.vega_mag bp sed) in\n Printf.printf \"%8s %+8.3f %+8.3f %+8.3f\\n\" name m_ab m_st m_vega)\n bands;\n\n Printf.printf \"\\n Note: AB and ST systems are defined by reference flux\\n\";\n Printf.printf \" densities; Vega magnitudes use the alpha Lyr spectrum.\\n\\n\";\n\n (* --- Part 2: K-corrections --- *)\n Printf.printf \"Part 2: K-corrections\\n\";\n Printf.printf \"---------------------\\n\\n\";\n\n let bp = Filters.sdss_r in\n let bp_wave = Photometry.wavelength bp in\n\n let rest_sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm |> Spectrum.as_flux_density\n in\n let m_ab_rest = item [] (Photometry.ab_mag bp rest_sed) in\n let m_st_rest = item [] (Photometry.st_mag bp rest_sed) in\n let m_vega_rest = item [] (Photometry.vega_mag bp rest_sed) in\n\n Printf.printf \" Rest-frame SDSS r-band:\\n\";\n Printf.printf \" AB = %.3f\\n\" m_ab_rest;\n Printf.printf \" ST = %.3f\\n\" m_st_rest;\n Printf.printf \" Vega = %.3f\\n\\n\" m_vega_rest;\n\n Printf.printf \" K-correction = m_obs(z) - m_rest\\n\\n\";\n Printf.printf \"%6s %8s %8s %8s\\n\" \"z\" \"K_AB\" \"K_ST\" \"K_Vega\";\n Printf.printf \"%6s %8s %8s %8s\\n\" \"------\" \"--------\" \"--------\" \"--------\";\n\n let redshifts = [| 0.1; 0.2; 0.3; 0.5; 0.7; 1.0 |] in\n Array.iter\n (fun z ->\n let zv = Nx.scalar f64 z in\n let obs_sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm |> Spectrum.as_flux_density\n |> Spectrum.redshift ~z:zv\n in\n let k_ab = item [] (Photometry.ab_mag bp obs_sed) -. m_ab_rest in\n let k_st = item [] (Photometry.st_mag bp obs_sed) -. m_st_rest in\n let k_vega = item [] (Photometry.vega_mag bp obs_sed) -. m_vega_rest in\n Printf.printf \"%6.2f %+8.3f %+8.3f %+8.3f\\n\" z k_ab k_st k_vega)\n redshifts;\n\n Printf.printf \"\\n\";\n\n (* --- Part 3: Color evolution with redshift --- *)\n Printf.printf \"Part 3: Color evolution (u-r) with redshift\\n\";\n Printf.printf \"-------------------------------------------\\n\\n\";\n Printf.printf \"%6s %8s\\n\" \"z\" \"u-r (AB)\";\n Printf.printf \"%6s %8s\\n\" \"------\" \"--------\";\n Array.iter\n (fun z ->\n let zv = Nx.scalar f64 z in\n let color =\n item []\n (Photometry.color Filters.sdss_u Filters.sdss_r\n (Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm |> Spectrum.as_flux_density\n |> Spectrum.redshift ~z:zv))\n in\n Printf.printf \"%6.2f %+8.3f\\n\" z color)\n redshifts;\n\n Printf.printf \"\\n\";\n\n (* --- Part 4: Extinction --- *)\n Printf.printf \"Part 4: Dust extinction\\n\";\n Printf.printf \"-----------------------\\n\\n\";\n\n let rv = Nx.scalar f64 3.1 in\n let av_values = [| 0.0; 0.5; 1.0; 2.0; 3.0 |] in\n\n Printf.printf \" CCM89 extinction law (R_V = 3.1)\\n\";\n Printf.printf \" Reddening a T=6000 K blackbody through SDSS r-band\\n\\n\";\n Printf.printf \"%6s %8s %8s %8s\\n\" \"A_V\" \"m_AB\" \"delta_m\" \"E(u-r)\";\n Printf.printf \"%6s %8s %8s %8s\\n\" \"------\" \"--------\" \"--------\" \"--------\";\n\n let unreddened_sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm |> Spectrum.as_flux_density\n in\n let m0 = item [] (Photometry.ab_mag bp unreddened_sed) in\n let color0 =\n item [] (Photometry.color Filters.sdss_u Filters.sdss_r unreddened_sed)\n in\n\n Array.iter\n (fun av_f ->\n let av = Nx.scalar f64 av_f in\n let reddened =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm\n |> Extinction.apply (Extinction.ccm89 ~rv) ~av\n |> Spectrum.as_flux_density\n in\n let m = item [] (Photometry.ab_mag bp reddened) in\n let color =\n item [] (Photometry.color Filters.sdss_u Filters.sdss_r reddened)\n in\n Printf.printf \"%6.1f %8.3f %+8.3f %+8.3f\\n\" av_f m (m -. m0)\n (color -. color0))\n av_values;\n\n Printf.printf \"\\n\";\n\n (* Compare extinction laws *)\n Printf.printf \" Comparing extinction laws at A_V = 1.0:\\n\\n\";\n Printf.printf \"%16s %8s %8s\\n\" \"Law\" \"r-band\" \"E(u-r)\";\n Printf.printf \"%16s %8s %8s\\n\" \"----------------\" \"--------\" \"--------\";\n\n let av_one = Nx.scalar f64 1.0 in\n let laws =\n [|\n (\"CCM89\", Extinction.ccm89 ~rv);\n (\"Fitzpatrick99\", Extinction.fitzpatrick99 ~rv);\n (\"O'Donnell94\", Extinction.odonnell94 ~rv);\n |]\n in\n Array.iter\n (fun (name, law) ->\n let reddened =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm\n |> Extinction.apply law ~av:av_one\n |> Spectrum.as_flux_density\n in\n let m = item [] (Photometry.ab_mag bp reddened) in\n let color =\n item [] (Photometry.color Filters.sdss_u Filters.sdss_r reddened)\n in\n Printf.printf \"%16s %+8.3f %+8.3f\\n\" name (m -. m0) (color -. color0))\n laws\n"
},
{
"path": "dev/umbra/examples/05-sed-fitting/README.md",
"content": "# `05-sed-fitting`\n\nFull SED fitting pipeline: fits stellar temperature, dust extinction (A_V), and\nflux normalization simultaneously from UGRIZ photometry. Demonstrates the\ncomposable differentiable pipeline through Spectrum, Extinction, and Photometry.\n\n```bash\ndune exec dev/umbra/examples/05-sed-fitting/main.exe\n```\n\n## What You'll Learn\n\n- Building a full astrophysical forward model from composable modules\n- How the blackbody -> extinction -> photometry pipeline is end-to-end differentiable\n- Creating custom bandpasses with `Photometry.tophat`\n- Fitting multiple correlated parameters (T, A_V, normalization) simultaneously\n\n## Key Functions\n\n| Function | Purpose |\n| ---------------------------- | --------------------------------------------- |\n| `Spectrum.blackbody` | Planck spectral radiance at given wavelengths |\n| `Spectrum.scale` | Scale spectrum values by a factor |\n| `Spectrum.as_flux_density` | Cast to flux density kind for photometry |\n| `Extinction.ccm89` | Create CCM89 extinction law with R_V |\n| `Extinction.apply` | Apply dust reddening to a spectrum |\n| `Photometry.tophat` | Create a rectangular bandpass |\n| `Photometry.ab_mag` | Compute AB magnitude through a bandpass |\n| `Rune.value_and_grads` | Autodiff through the entire pipeline |\n\n## How It Works\n\nThe forward model constructs a synthetic SED at each optimization step:\n\n1. **Spectrum.blackbody** generates the Planck function at temperature T\n2. **Spectrum.scale** applies the flux normalization\n3. **Extinction.apply** reddens the spectrum using CCM89 with extinction A_V\n4. **Photometry.ab_mag** integrates through each bandpass to produce magnitudes\n\nSince every step is built from Nx tensor operations, Rune computes gradients\nof chi-squared with respect to all three parameters (log T, A_V, log norm) in\na single backward pass.\n\nThe temperature and normalization are parameterized in log-space for positivity\nand better gradient conditioning. A_V is left in linear space since it can\nmeaningfully be zero or negative (de-reddening).\n\n## Try It\n\n1. Replace tophat filters with real SDSS filters from `Filters.sdss_u`, etc.\n2. Add a redshift parameter to fit photometric redshifts.\n3. Try `Extinction.fitzpatrick99` instead of `ccm89` and compare results.\n4. Increase the photometric noise and observe how parameter uncertainties grow.\n\n## Next Steps\n\nContinue to [06-coordinates-and-time](../06-coordinates-and-time/) to work with\ncelestial coordinates, time scales, and observing conditions.\n"
},
{
"path": "dev/umbra/examples/05-sed-fitting/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega umbra))\n"
},
{
"path": "dev/umbra/examples/05-sed-fitting/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Differentiable SED fitting: temperature + extinction from photometry.\n\n Demonstrates the composable differentiable pipeline: Spectrum.blackbody ->\n Extinction.apply -> Photometry.ab_mag\n\n All operations flow through Nx tensor ops, making the entire pipeline\n differentiable via Rune's autodiff. We fit stellar temperature, dust\n extinction, and flux normalization simultaneously by gradient descent on\n photometric residuals. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\n(* Define 5 broadband filters (UGRIZ-like tophats) *)\nlet n_bp = 100\n\nlet band_u =\n Photometry.tophat ~lo:(Unit.Length.m 3.0e-7) ~hi:(Unit.Length.m 4.0e-7)\n ~n:n_bp\n\nlet band_g =\n Photometry.tophat ~lo:(Unit.Length.m 4.0e-7) ~hi:(Unit.Length.m 5.5e-7)\n ~n:n_bp\n\nlet band_r =\n Photometry.tophat ~lo:(Unit.Length.m 5.5e-7) ~hi:(Unit.Length.m 7.0e-7)\n ~n:n_bp\n\nlet band_i =\n Photometry.tophat ~lo:(Unit.Length.m 7.0e-7) ~hi:(Unit.Length.m 8.5e-7)\n ~n:n_bp\n\nlet band_z =\n Photometry.tophat ~lo:(Unit.Length.m 8.5e-7) ~hi:(Unit.Length.m 1.0e-6)\n ~n:n_bp\n\nlet bands = [ band_u; band_g; band_r; band_i; band_z ]\nlet band_names = [| \"U\"; \"G\"; \"R\"; \"I\"; \"Z\" |]\n\n(* True parameters *)\nlet true_temp = 6500.0 (* K -- F-type star *)\nlet true_av = 0.5 (* moderate extinction *)\nlet true_log_norm = -50.0\n\n(* Generate synthetic observations *)\nlet rv = Nx.scalar f64 3.1\n\nlet obs_mags =\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 true_temp) in\n let av = Nx.scalar f64 true_av in\n let norm = Nx.scalar f64 (Float.exp true_log_norm) in\n let mags =\n List.map\n (fun bp ->\n let bp_wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm\n |> Extinction.apply (Extinction.ccm89 ~rv) ~av\n |> Spectrum.as_flux_density\n in\n Photometry.ab_mag bp sed)\n bands\n in\n (* Add 3% photometric noise *)\n let noise = [| 0.03; -0.02; 0.01; -0.01; 0.02 |] in\n List.mapi (fun i m -> Nx.add_s m noise.(i)) mags\n\nlet obs_errs = List.init 5 (fun _ -> Nx.scalar f64 0.05)\n\n(* Forward model: generate magnitudes from parameters *)\nlet forward_model log_temp av log_norm =\n let temp = Unit.Temperature.of_kelvin (exp log_temp) in\n let norm = exp log_norm in\n List.map\n (fun bp ->\n let bp_wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm\n |> Extinction.apply (Extinction.ccm89 ~rv) ~av\n |> Spectrum.as_flux_density\n in\n Photometry.ab_mag bp sed)\n bands\n\n(* Loss function: chi-squared *)\nlet loss params =\n match params with\n | [ log_temp; av; log_norm ] ->\n let pred = forward_model log_temp av log_norm in\n List.fold_left2\n (fun acc p (o, e) ->\n let residual = div (sub p o) e in\n add acc (square residual))\n (scalar f64 0.0) pred\n (List.combine obs_mags obs_errs)\n | _ -> failwith \"expected [log_temp; av; log_norm]\"\n\nlet () =\n Printf.printf \"Differentiable SED Fitting\\n\";\n Printf.printf \"=========================\\n\";\n Printf.printf\n \"Pipeline: Spectrum.blackbody -> Extinction.ccm89 -> Photometry.ab_mag\\n\\n\";\n\n Printf.printf \"True parameters:\\n\";\n Printf.printf \" T = %.0f K\\n\" true_temp;\n Printf.printf \" A_V = %.2f mag\\n\" true_av;\n Printf.printf \" logN = %.1f\\n\\n\" true_log_norm;\n\n Printf.printf \"Observed magnitudes (with noise):\\n\";\n List.iteri\n (fun i m ->\n Printf.printf \" %s = %.3f +/- %.3f\\n\" band_names.(i) (item [] m)\n (item [] (List.nth obs_errs i)))\n obs_mags;\n Printf.printf \"\\n\";\n\n (* Initial guesses *)\n let algo = Vega.adam (Vega.Schedule.constant 1e-3) in\n let log_temp = ref (scalar f64 (Float.log 7000.0)) in\n let av = ref (scalar f64 0.3) in\n let log_norm = ref (scalar f64 (-50.5)) in\n let states =\n [| Vega.init algo !log_temp; Vega.init algo !av; Vega.init algo !log_norm |]\n in\n let steps = 1000 in\n\n Printf.printf \"%5s %10s %8s %8s %8s\\n\" \"step\" \"chi2\" \"T (K)\" \"A_V\"\n \"log_norm\";\n Printf.printf \"%5s %10s %8s %8s %8s\\n\" \"-----\" \"----------\" \"--------\"\n \"--------\" \"--------\";\n\n let refs = [| log_temp; av; log_norm |] in\n for i = 0 to steps - 1 do\n let loss_val, grads =\n Rune.value_and_grads loss [ !log_temp; !av; !log_norm ]\n in\n if i mod 200 = 0 || i = steps - 1 then\n Printf.printf \"%5d %10.4f %8.1f %8.3f %8.3f\\n\" i (item [] loss_val)\n (Float.exp (item [] !log_temp))\n (item [] !av) (item [] !log_norm);\n List.iteri\n (fun j g ->\n let p, s = Vega.step states.(j) ~grad:g ~param:!(refs.(j)) in\n refs.(j) := p;\n states.(j) <- s)\n grads\n done;\n\n Printf.printf \"\\nFitted parameters:\\n\";\n Printf.printf \" T = %.1f K (true: %.0f K)\\n\"\n (Float.exp (item [] !log_temp))\n true_temp;\n Printf.printf \" A_V = %.3f (true: %.2f)\\n\" (item [] !av) true_av;\n Printf.printf \" logN = %.3f (true: %.1f)\\n\" (item [] !log_norm)\n true_log_norm;\n\n (* Show fitted vs observed magnitudes *)\n Printf.printf \"\\nFitted vs observed magnitudes:\\n\";\n let fitted_mags = forward_model !log_temp !av !log_norm in\n Printf.printf \"%5s %8s %8s %8s\\n\" \"Band\" \"Observed\" \"Fitted\" \"Residual\";\n Printf.printf \"%5s %8s %8s %8s\\n\" \"-----\" \"--------\" \"--------\" \"--------\";\n List.iteri\n (fun i (obs, fit) ->\n let o = item [] obs in\n let f = item [] fit in\n Printf.printf \"%5s %8.3f %8.3f %+8.3f\\n\" band_names.(i) o f (f -. o))\n (List.combine obs_mags fitted_mags)\n"
},
{
"path": "dev/umbra/examples/06-coordinates-and-time/README.md",
"content": "# `06-coordinates-and-time`\n\nCelestial coordinates, astronomical time scales, and survey selection.\nDemonstrates frame transforms (ICRS, Galactic), angular separation, time scale\nconversions (UTC, TAI, TT, TDB), altitude-azimuth coordinates, airmass, and\na practical survey selection function.\n\n```bash\ndune exec dev/umbra/examples/06-coordinates-and-time/main.exe\n```\n\n## What You'll Learn\n\n- Creating celestial coordinates in ICRS and converting to Galactic frame\n- Computing angular separations between objects\n- Parsing ISO 8601 dates and converting between time scales\n- Computing horizontal coordinates for a ground-based observer\n- Building a survey selection function from airmass, altitude, and magnitude cuts\n\n## Key Functions\n\n| Function | Purpose |\n| ---------------------------- | --------------------------------------------- |\n| `Coord.of_radec` | Create ICRS coordinates from RA/Dec |\n| `Coord.galactic` | Convert to Galactic coordinates |\n| `Coord.separation` | Angular separation between positions |\n| `Time.of_iso` | Parse ISO 8601 date-time as UTC |\n| `Time.utc_to_tai` | Convert UTC to TAI |\n| `Time.tai_to_tt` | Convert TAI to Terrestrial Time |\n| `Time.tt_to_tdb` | Convert TT to Barycentric Dynamical Time |\n| `Time.to_jd`, `Time.to_mjd` | Extract Julian Date / Modified Julian Date |\n| `Altaz.make_observer` | Create a ground-based observer location |\n| `Altaz.of_coord` | Convert celestial to horizontal coordinates |\n| `Altaz.alt`, `Altaz.az` | Extract altitude and azimuth |\n| `Altaz.airmass` | Compute airmass at given altitude |\n| `Filters.rubin_r` | Pre-built Rubin/LSST r-band filter |\n\n## How It Works\n\n**Coordinates**: Positions are stored as (longitude, latitude) pairs in typed\nangle quantities. Frame transforms use 3x3 rotation matrices to convert\nbetween ICRS, Galactic, Ecliptic, and Supergalactic systems. Angular separation\nuses the Vincenty formula for numerical stability.\n\n**Time**: Julian Dates carry phantom type tags (UTC, TAI, TT, TDB) that\nenforce correct scale conversions at compile time. UTC-TAI uses the IERS\nleap-second table; TT = TAI + 32.184s exactly; TDB-TT uses the Fairhead &\nBretagnon series.\n\n**Altaz**: Converts ICRS to horizontal coordinates using IAU 2006 precession\nand the Earth Rotation Angle. Airmass uses the Pickering (2002) formula.\n\n**Selection**: Combines altitude (above horizon), airmass (atmospheric\nextinction), and magnitude limit into a boolean selection function -- a\nbuilding block for survey simulations.\n\n## Try It\n\n1. Add atmospheric refraction with `Altaz.of_coord ~refraction:true`.\n2. Compute the position angle from Vega to Deneb with `Coord.position_angle`.\n3. Use `Coord.of_galactic` to create coordinates in the Galactic plane and\n convert to ICRS.\n4. Change the observer location and time to see how visibility changes.\n\n## Next Steps\n\nExplore the other Umbra examples for more advanced topics: catalog\ncross-matching with `Coord.nearest`, cosmological power spectra with\n`Cosmo.linear_power`, and Fisher matrix forecasts.\n"
},
{
"path": "dev/umbra/examples/06-coordinates-and-time/dune",
"content": "(executable\n (name main)\n (libraries nx rune umbra))\n"
},
{
"path": "dev/umbra/examples/06-coordinates-and-time/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Coordinates, time scales, and survey selection.\n\n Demonstrates Umbra's coordinate, time, and observing modules: - Coord:\n celestial coordinates with frame transforms (ICRS, Galactic, Ecliptic,\n Supergalactic) and angular separation. - Time: astronomical time with\n type-safe scale conversions (UTC, TAI, TT, TDB) and ISO 8601 parsing. -\n Altaz: horizontal coordinates, airmass, and atmospheric refraction.\n\n Combines these into a survey selection function that determines which targets\n are observable given an observer, time, and observing constraints. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\nlet () =\n Printf.printf \"Coordinates, time scales, and survey selection\\n\";\n Printf.printf \"===============================================\\n\\n\";\n\n (* --- Part 1: Coordinate frames --- *)\n Printf.printf \"Part 1: Coordinate frame transforms\\n\";\n Printf.printf \"------------------------------------\\n\\n\";\n\n let targets =\n [|\n (\"Galactic center\", 266.417, -28.936);\n (\"Vega\", 279.235, 38.784);\n (\"North Galactic Pole\", 192.860, 27.128);\n (\"LMC\", 80.894, -69.756);\n (\"M31 (Andromeda)\", 10.685, 41.269);\n |]\n in\n\n Printf.printf \"%20s %8s %8s %8s %8s\\n\" \"Object\" \"RA\" \"Dec\" \"l\" \"b\";\n Printf.printf \"%20s %8s %8s %8s %8s\\n\" \"--------------------\" \"--------\"\n \"--------\" \"--------\" \"--------\";\n\n Array.iter\n (fun (name, ra_deg, dec_deg) ->\n let coord =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| ra_deg |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| dec_deg |]))\n in\n let gal = Coord.galactic coord in\n let l = item [ 0 ] (Unit.Angle.in_deg (Coord.lon gal)) in\n let b = item [ 0 ] (Unit.Angle.in_deg (Coord.lat gal)) in\n Printf.printf \"%20s %8.2f %+8.2f %8.2f %+8.2f\\n\" name ra_deg dec_deg l\n b)\n targets;\n Printf.printf \"\\n\";\n\n (* Angular separation *)\n Printf.printf \"Angular separations:\\n\";\n let vega =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 279.235 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 38.784 |]))\n in\n let altair =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 297.696 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 8.868 |]))\n in\n let deneb =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 310.358 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 45.280 |]))\n in\n let sep_va = item [ 0 ] (Unit.Angle.in_deg (Coord.separation vega altair)) in\n let sep_vd = item [ 0 ] (Unit.Angle.in_deg (Coord.separation vega deneb)) in\n let sep_ad = item [ 0 ] (Unit.Angle.in_deg (Coord.separation altair deneb)) in\n Printf.printf \" Vega - Altair: %.2f deg\\n\" sep_va;\n Printf.printf \" Vega - Deneb: %.2f deg\\n\" sep_vd;\n Printf.printf \" Altair - Deneb: %.2f deg\\n\" sep_ad;\n Printf.printf \" (Summer Triangle)\\n\\n\";\n\n (* --- Part 2: Time scales --- *)\n Printf.printf \"Part 2: Astronomical time scales\\n\";\n Printf.printf \"--------------------------------\\n\\n\";\n\n let t_utc = Time.of_iso \"2024-06-21T04:00:00\" in\n let t_tai = Time.utc_to_tai t_utc in\n let t_tt = Time.tai_to_tt t_tai in\n let t_tdb = Time.tt_to_tdb t_tt in\n\n Printf.printf \" UTC: %s\\n\" (Time.to_iso t_utc);\n Printf.printf \" JD (UTC): %.6f\\n\" (Time.to_jd t_utc);\n Printf.printf \" MJD (UTC): %.6f\\n\" (Time.to_mjd t_utc);\n Printf.printf \" JD (TAI): %.6f\\n\" (Time.to_jd t_tai);\n Printf.printf \" JD (TT): %.6f\\n\" (Time.to_jd t_tt);\n Printf.printf \" JD (TDB): %.6f\\n\" (Time.to_jd t_tdb);\n\n let dt_tai_utc =\n Unit.to_float (Time.diff t_tai (Time.unsafe_of_jd (Time.to_jd t_utc)))\n in\n Printf.printf \"\\n TAI - UTC = %.1f s (leap seconds)\\n\" (dt_tai_utc *. 86400.0);\n\n let t_j2000 = Time.of_iso \"2000-01-01T12:00:00\" in\n let dt_j2000 = Unit.to_float (Time.diff t_utc t_j2000) in\n Printf.printf \" Days since J2000.0: %.2f\\n\\n\" (dt_j2000 *. 86400.0 /. 86400.0);\n\n (* --- Part 3: Horizontal coordinates and airmass --- *)\n Printf.printf \"Part 3: Altitude-azimuth and airmass\\n\";\n Printf.printf \"------------------------------------\\n\\n\";\n\n (* Observer at Cerro Pachon (Rubin site) *)\n let obs =\n Altaz.make_observer\n ~lat:(Unit.Angle.deg (-30.2444))\n ~lon:(Unit.Angle.deg (-70.7494))\n ~height:(Unit.Length.m 2663.0) ()\n in\n let obstime = Time.of_iso \"2024-06-21T04:00:00\" in\n\n Printf.printf \" Observer: Cerro Pachon (Rubin Observatory)\\n\";\n Printf.printf \" Lat: %.4f deg\\n\" (-30.2444);\n Printf.printf \" Lon: %.4f deg\\n\" (-70.7494);\n Printf.printf \" Elevation: %.0f m\\n\" 2663.0;\n Printf.printf \" Time: 2024-06-21 04:00 UTC\\n\\n\";\n\n let stars =\n [|\n (\"Vega\", 279.235, 38.784);\n (\"Sirius\", 101.287, -16.716);\n (\"Canopus\", 95.988, -52.696);\n (\"Alpha Cen\", 219.902, -60.834);\n (\"Fomalhaut\", 344.413, -29.622);\n |]\n in\n\n Printf.printf \"%12s %7s %7s %8s\\n\" \"Star\" \"Alt\" \"Az\" \"Airmass\";\n Printf.printf \"%12s %7s %7s %8s\\n\" \"------------\" \"-------\" \"-------\"\n \"--------\";\n\n Array.iter\n (fun (name, ra_deg, dec_deg) ->\n let coord =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| ra_deg |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| dec_deg |]))\n in\n let hz = Altaz.of_coord ~obstime ~observer:obs coord in\n let alt_deg =\n item [ 0 ] (Unit.Angle.to_tensor (Altaz.alt hz)) *. 180.0 /. Float.pi\n in\n let az_deg =\n item [ 0 ] (Unit.Angle.to_tensor (Altaz.az hz)) *. 180.0 /. Float.pi\n in\n let am = item [ 0 ] (Altaz.airmass hz) in\n Printf.printf \"%12s %+7.1f %7.1f %8.2f\\n\" name alt_deg az_deg am)\n stars;\n Printf.printf \"\\n\";\n\n (* --- Part 4: Survey selection --- *)\n Printf.printf \"Part 4: Survey selection function\\n\";\n Printf.printf \"---------------------------------\\n\\n\";\n\n let mag_limit = 20.0 in\n let airmass_cut = 2.0 in\n Printf.printf \" Selection criteria:\\n\";\n Printf.printf \" Magnitude limit: r < %.1f (AB)\\n\" mag_limit;\n Printf.printf \" Airmass cut: X < %.1f\\n\" airmass_cut;\n Printf.printf \" Above horizon: alt > 0 deg\\n\\n\";\n\n let bp = Filters.rubin_r in\n let norm = Nx.scalar f64 (Float.exp (-49.0)) in\n\n let star_data =\n [|\n (\"Vega\", 279.235, 38.784, 5800.0);\n (\"Sirius\", 101.287, -16.716, 9940.0);\n (\"Canopus\", 95.988, -52.696, 7350.0);\n (\"Alpha Cen\", 219.902, -60.834, 5790.0);\n (\"Fomalhaut\", 344.413, -29.622, 8590.0);\n |]\n in\n\n Printf.printf \"%12s %7s %8s %6s %s\\n\" \"Star\" \"Alt\" \"Airmass\" \"r_mag\"\n \"Select?\";\n Printf.printf \"%12s %7s %8s %6s %s\\n\" \"------------\" \"-------\" \"--------\"\n \"------\" \"-------\";\n\n Array.iter\n (fun (name, ra_deg, dec_deg, temp_k) ->\n let coord =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| ra_deg |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| dec_deg |]))\n in\n let hz = Altaz.of_coord ~obstime ~observer:obs coord in\n let alt_deg =\n item [ 0 ] (Unit.Angle.to_tensor (Altaz.alt hz)) *. 180.0 /. Float.pi\n in\n let am = item [ 0 ] (Altaz.airmass hz) in\n\n (* Synthetic magnitude through Rubin r-band *)\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 temp_k) in\n let bp_wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm |> Spectrum.as_flux_density\n in\n let r_mag = item [] (Photometry.ab_mag bp sed) in\n\n let selected = alt_deg > 0.0 && am < airmass_cut && r_mag < mag_limit in\n Printf.printf \"%12s %+7.1f %8.2f %6.2f %s\\n\" name alt_deg am r_mag\n (if selected then \"YES\" else \"no\"))\n star_data;\n\n Printf.printf \"\\n Height stored: %.0f m\\n\"\n (item [] (Unit.Length.to_tensor (Altaz.observer_height obs)))\n"
},
{
"path": "dev/umbra/examples/07-batch-photometry/README.md",
"content": "# `07-batch-photometry`\n\nComputes SDSS g-r colors for a grid of blackbody templates at different\ntemperatures and dust extinctions in a single pass using batch operations.\nInstead of looping over individual spectra, the values tensor has a leading\nbatch dimension and all photometry operations broadcast over it.\n\n```bash\ncd dev/umbra\ndune exec --root . examples/07-batch-photometry/main.exe\n```\n\n## What You'll Learn\n\n- Constructing batched spectra by stacking blackbodies into a leading dimension\n- Broadcasting extinction across a batch of SEDs with per-spectrum A_V\n- Computing synthetic SDSS photometry with AB magnitudes\n- Exploring color-temperature and color-extinction relations\n\n## Key Functions\n\n| Function | Purpose |\n| -------------------------- | ------------------------------------------------ |\n| `Spectrum.blackbody` | Generate Planck spectrum at a given temperature |\n| `Spectrum.create` | Build a spectrum from wavelength and value arrays |\n| `Spectrum.as_flux_density` | Cast to flux density kind for photometry |\n| `Nx.stack` | Stack individual spectra into a batch dimension |\n| `Extinction.ccm89` | Create CCM89 dust extinction law |\n| `Extinction.apply` | Apply reddening with per-spectrum A_V broadcast |\n| `Photometry.ab_mag` | Compute AB magnitude through a bandpass |\n| `Filters.sdss_g` | SDSS g-band filter response |\n\n## How It Works\n\nThe example first builds a grid of 20 blackbody spectra from 3000 K to 30000 K\nby stacking individual `Spectrum.blackbody` outputs into a `[n_temp; 500]`\nvalues tensor. When this batch spectrum is passed to `Photometry.ab_mag`, the\nintegration broadcasts over the leading dimension, producing one magnitude per\ntemperature in a single call.\n\nThe second half demonstrates per-spectrum extinction. A T=6000 K blackbody is\nreplicated into 10 copies, and `Extinction.apply` is called with an A_V tensor\nof shape `[n_av; 1]` that broadcasts against the `[n_av; 500]` flux values.\nThis yields reddened g-r colors across a range of dust columns without any\nexplicit loop.\n\n## Try It\n\n1. Increase the temperature grid to 100 points and plot the g-r color curve to\n see where the blue turnover occurs.\n2. Add a third band (sdss_i) and compute the g-r vs r-i color-color diagram.\n3. Replace the blackbody with a power-law spectrum and observe how the color\n trends differ.\n\n## Next Steps\n\nContinue to [08-photometric-redshifts](../08-photometric-redshifts/) to learn\nhow to estimate galaxy redshifts by combining grid search with gradient-based\nrefinement through the differentiable photometry pipeline.\n"
},
{
"path": "dev/umbra/examples/07-batch-photometry/dune",
"content": "(executable\n (name main)\n (libraries nx rune umbra))\n"
},
{
"path": "dev/umbra/examples/07-batch-photometry/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Batch template photometry.\n\n Computes SDSS g-r colors for a grid of blackbody templates at different\n temperatures and dust extinctions in a single pass, demonstrating batched\n spectra. Instead of looping over individual spectra, the values tensor has a\n leading batch dimension and all photometry operations broadcast over it. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\nlet () =\n Printf.printf \"Batch Template Photometry\\n\";\n Printf.printf \"=========================\\n\\n\";\n\n (* Temperature grid: 20 blackbodies from 3000K to 30000K *)\n let n_temp = 20 in\n let temps =\n Array.init n_temp (fun i ->\n 3000.0\n +. (Float.of_int i *. (30000.0 -. 3000.0) /. Float.of_int (n_temp - 1)))\n in\n\n (* Shared wavelength grid covering SDSS g and r *)\n let wavelength = Unit.Length.of_m (Nx.linspace f64 3e-7 1.1e-6 500) in\n\n (* Build batch spectrum: stack individual blackbodies into [n_temp; 500] *)\n let values =\n Nx.stack\n (List.init n_temp (fun i ->\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 temps.(i)) in\n Spectrum.values (Spectrum.blackbody ~temperature:temp ~wavelength)))\n in\n let batch = Spectrum.create ~wavelength ~values |> Spectrum.as_flux_density in\n\n (* AB magnitudes in g and r — returns shape [n_temp] each *)\n let g_mag = Photometry.ab_mag Filters.sdss_g batch in\n let r_mag = Photometry.ab_mag Filters.sdss_r batch in\n let g_r = Nx.sub g_mag r_mag in\n\n Printf.printf \"Unreddened blackbody colors (SDSS g-r):\\n\";\n Printf.printf \"%8s %8s %8s %8s\\n\" \"T (K)\" \"g\" \"r\" \"g-r\";\n Printf.printf \"%8s %8s %8s %8s\\n\" \"--------\" \"--------\" \"--------\"\n \"--------\";\n Array.iteri\n (fun i t ->\n if i mod 4 = 0 || i = n_temp - 1 then\n Printf.printf \"%8.0f %+8.3f %+8.3f %+8.3f\\n\" t (item [ i ] g_mag)\n (item [ i ] r_mag) (item [ i ] g_r))\n temps;\n\n (* Now apply per-spectrum extinction: A_V from 0.0 to 2.0 *)\n Printf.printf \"\\nReddening a T=6000K blackbody (SDSS g-r vs A_V):\\n\";\n let n_av = 10 in\n let av_values = Nx.linspace f64 0.0 2.0 n_av in\n\n (* Single-temperature spectrum, batched over A_V *)\n let temp_6k = Unit.Temperature.of_kelvin (Nx.scalar f64 6000.0) in\n let sed_1d =\n Spectrum.blackbody ~temperature:temp_6k ~wavelength\n |> Spectrum.as_flux_density\n in\n (* Replicate into [n_av; 500] *)\n let sed_values =\n Nx.stack (List.init n_av (fun _ -> Spectrum.values sed_1d))\n in\n let sed_batch =\n Spectrum.create ~wavelength ~values:sed_values |> Spectrum.as_flux_density\n in\n (* Per-spectrum A_V: reshape to [n_av; 1] to broadcast with [n_av; 500] *)\n let rv = Nx.scalar f64 3.1 in\n let av_col = Nx.reshape [| n_av; 1 |] av_values in\n let reddened = Extinction.apply (Extinction.ccm89 ~rv) ~av:av_col sed_batch in\n let g_red = Photometry.ab_mag Filters.sdss_g reddened in\n let r_red = Photometry.ab_mag Filters.sdss_r reddened in\n let g_r_red = Nx.sub g_red r_red in\n\n Printf.printf \"%8s %8s\\n\" \"A_V\" \"g-r\";\n Printf.printf \"%8s %8s\\n\" \"--------\" \"--------\";\n for i = 0 to n_av - 1 do\n Printf.printf \"%8.2f %+8.3f\\n\" (item [ i ] av_values) (item [ i ] g_r_red)\n done\n"
},
{
"path": "dev/umbra/examples/08-photometric-redshifts/README.md",
"content": "# `08-photometric-redshifts`\n\nTwo-stage photometric redshift estimation: coarse grid search followed by\ngradient-based refinement using Adam. The full pipeline (blackbody -> redshift\n-> extinction -> photometry) is differentiable through Rune, enabling gradient\ndescent on redshift and normalization parameters against synthetic SDSS ugriz\nobservations.\n\n```bash\ncd dev/umbra\ndune exec --root . examples/08-photometric-redshifts/main.exe\n```\n\n## What You'll Learn\n\n- Building an end-to-end differentiable photometric pipeline through SDSS ugriz filters\n- Composing spectrum redshifting, dust extinction, and synthetic photometry\n- Combining grid search initialization with autodiff gradient refinement\n- Using multi-parameter gradients to jointly fit redshift and normalization\n\n## Key Functions\n\n| Function | Purpose |\n| -------------------------- | ---------------------------------------------------- |\n| `Spectrum.blackbody` | Generate a template SED at given temperature |\n| `Spectrum.redshift` | Apply cosmological redshift to a spectrum |\n| `Spectrum.scale` | Scale spectrum by a normalization factor |\n| `Extinction.apply` | Apply dust reddening with CCM89 law |\n| `Photometry.ab_mag` | Compute AB magnitude through a bandpass |\n| `Photometry.wavelength` | Extract the wavelength grid of a bandpass filter |\n| `Rune.value_and_grads` | Compute loss and parameter gradients in one pass |\n| `Vega.adam` | Adam optimizer for gradient refinement |\n\n## How It Works\n\nThe example generates synthetic observed magnitudes for a galaxy at z=0.3 with\nT=5500 K, A_V=0.2, by pushing a blackbody through the full pipeline:\n`blackbody -> scale -> extinction -> redshift -> ab_mag` in each of the five\nSDSS bands. These serve as the \"data\" to fit against.\n\nStage 1 performs a coarse grid search over 30 redshift values from 0.01 to\n0.90, computing chi-squared at each point with a fixed template. This\nidentifies a rough minimum without requiring gradients.\n\nStage 2 takes the best grid redshift and refines it with 500 Adam optimizer\nsteps. The loss function (sum of squared magnitude residuals) flows through\n`Spectrum.redshift` and `Photometry.ab_mag`, so Rune provides exact gradients\nwith respect to log(1+z) and log(normalization). The parameterization in\nlog-space ensures positivity and improves conditioning.\n\n## Try It\n\n1. Change the true redshift to z=0.7 and observe how the grid search coarseness\n affects the initial estimate.\n2. Add temperature as a third free parameter in the refinement stage.\n3. Replace the single blackbody template with a composite SED that includes an\n emission line.\n\n## Next Steps\n\nContinue to [09-gravitational-lensing](../09-gravitational-lensing/) to see how\nRune's autodiff can fit physical parameters of a gravitational lens model from\nobserved image positions.\n"
},
{
"path": "dev/umbra/examples/08-photometric-redshifts/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega umbra))\n"
},
{
"path": "dev/umbra/examples/08-photometric-redshifts/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Photometric redshift estimation via template fitting.\n\n Demonstrates composing Spectrum.redshift -> Extinction.apply ->\n Photometry.ab_mag through real SDSS filters, with gradient refinement via\n Rune's autodiff. Auto-resampling makes the pipeline seamless.\n\n Stage 1: Grid search over redshift to find a coarse estimate. Stage 2: Adam\n optimizer refines z and normalization using AD gradients. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\nlet bands =\n [\n Filters.sdss_u;\n Filters.sdss_g;\n Filters.sdss_r;\n Filters.sdss_i;\n Filters.sdss_z;\n ]\n\nlet band_names = [| \"u\"; \"g\"; \"r\"; \"i\"; \"z\" |]\n\n(* True parameters for synthetic galaxy *)\nlet true_z = 0.3\nlet true_temp = 5500.0\nlet true_av = 0.2\nlet true_log_norm = -50.0\nlet rv = Nx.scalar f64 3.1\n\n(* Synthetic observed magnitudes *)\nlet obs_mags =\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 true_temp) in\n let z = Nx.scalar f64 true_z in\n let av = Nx.scalar f64 true_av in\n let norm = Nx.scalar f64 (Float.exp true_log_norm) in\n List.map\n (fun bp ->\n let bp_wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm\n |> Extinction.apply (Extinction.ccm89 ~rv) ~av\n |> Spectrum.as_flux_density |> Spectrum.redshift ~z\n in\n Photometry.ab_mag bp sed)\n bands\n\n(* Grid search: coarse scan over z *)\nlet grid_search () =\n let best_z = ref 0.0 in\n let best_chi2 = ref Float.infinity in\n let n_z = 30 in\n for iz = 0 to n_z - 1 do\n let z = Nx.scalar f64 (0.01 +. (Float.of_int iz *. 0.03)) in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 5000.0) in\n let norm = Nx.scalar f64 (Float.exp (-50.0)) in\n let pred =\n List.map\n (fun bp ->\n let bp_wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm |> Spectrum.as_flux_density\n |> Spectrum.redshift ~z\n in\n Photometry.ab_mag bp sed)\n bands\n in\n (* Color-based chi-squared: compare color differences *)\n let chi2 =\n List.fold_left2\n (fun acc p o -> add acc (square (sub p o)))\n (scalar f64 0.0) pred obs_mags\n in\n let chi2_v = item [] chi2 in\n if chi2_v < !best_chi2 then begin\n best_chi2 := chi2_v;\n best_z := item [] z\n end\n done;\n !best_z\n\n(* Gradient refinement around grid minimum *)\nlet refine z0 =\n let loss params =\n match params with\n | [ log_z1; log_norm ] ->\n let z = sub (exp log_z1) (scalar f64 1.0) in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 5500.0) in\n let norm = exp log_norm in\n let pred =\n List.map\n (fun bp ->\n let bp_wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:bp_wave\n |> Spectrum.scale norm |> Spectrum.as_flux_density\n |> Spectrum.redshift ~z\n in\n Photometry.ab_mag bp sed)\n bands\n in\n List.fold_left2\n (fun acc p o -> add acc (square (sub p o)))\n (scalar f64 0.0) pred obs_mags\n | _ -> failwith \"expected [log_z1; log_norm]\"\n in\n let algo = Vega.adam (Vega.Schedule.constant 5e-4) in\n let log_z1 = ref (scalar f64 (Float.log (1.0 +. z0))) in\n let log_norm = ref (scalar f64 (-50.0)) in\n let states = [| Vega.init algo !log_z1; Vega.init algo !log_norm |] in\n let refs = [| log_z1; log_norm |] in\n for _ = 0 to 499 do\n let _loss_val, grads = Rune.value_and_grads loss [ !log_z1; !log_norm ] in\n List.iteri\n (fun j g ->\n let p, s = Vega.step states.(j) ~grad:g ~param:!(refs.(j)) in\n refs.(j) := p;\n states.(j) <- s)\n grads\n done;\n Float.exp (item [] !log_z1) -. 1.0\n\nlet () =\n Printf.printf \"Photometric Redshift Estimation\\n\";\n Printf.printf \"===============================\\n\";\n Printf.printf\n \"Pipeline: blackbody -> redshift -> extinction -> ab_mag (SDSS)\\n\\n\";\n Printf.printf \"True: z=%.3f T=%.0fK A_V=%.2f\\n\\n\" true_z true_temp true_av;\n Printf.printf \"Observed magnitudes:\\n\";\n List.iteri\n (fun i m -> Printf.printf \" %s = %.3f\\n\" band_names.(i) (item [] m))\n obs_mags;\n Printf.printf \"\\nStep 1: Grid search (z = 0.01 to 0.90)...\\n\";\n let z_grid = grid_search () in\n Printf.printf \" Best grid z = %.3f\\n\" z_grid;\n Printf.printf \"\\nStep 2: Gradient refinement (500 Adam steps)...\\n\";\n let z_fit = refine z_grid in\n Printf.printf \" Refined z = %.4f (true: %.3f)\\n\" z_fit true_z;\n Printf.printf \" Error = %.4f\\n\" (Float.abs (z_fit -. true_z))\n"
},
{
"path": "dev/umbra/examples/09-gravitational-lensing/README.md",
"content": "# `09-gravitational-lensing`\n\nFits gravitational lens parameters (lens center and Einstein radius) from\nobserved image positions of a quadruply-imaged quasar. The point-mass lens\nequation is expressed as Nx tensor operations, making the model fully\ndifferentiable through Rune for gradient-based fitting with Adam.\n\n```bash\ncd dev/umbra\ndune exec --root . examples/09-gravitational-lensing/main.exe\n```\n\n## What You'll Learn\n\n- Expressing the gravitational lens equation as differentiable tensor operations\n- Minimizing source-plane variance to fit lens parameters\n- Fitting physical parameters (lens position, Einstein radius) via Adam optimizer\n- Using autodiff gradients with a physics-based loss function\n\n## Key Functions\n\n| Function | Purpose |\n| ----------------------- | ------------------------------------------------------ |\n| `Nx.square` | Squared distances for radial computation |\n| `Nx.sqrt` | Radial distance from lens center |\n| `Nx.mean` | Mean source position across images |\n| `Rune.value_and_grads` | Compute loss and gradients for all lens parameters |\n| `Vega.adam` | Adam optimizer for parameter fitting |\n| `Vega.step` | Apply one optimization update |\n\n## How It Works\n\nA point-mass gravitational lens deflects light according to the lens equation:\nbeta = theta - theta_E^2 * theta_hat / |theta|, where beta is the true source\nposition, theta is the observed image position, and theta_E is the Einstein\nradius. If the lens model is correct, all observed images should map back to the\nsame source position in the source plane.\n\nThe example generates synthetic image positions for a quadruply-imaged quasar\nwith known lens parameters (x_L=0.1, y_L=-0.05, theta_E=1.0) plus small noise.\nThe loss function maps each image back to the source plane using the current\nlens parameters and computes the variance of the inferred source positions. A\ncorrect lens model yields zero variance.\n\nStarting from an initial guess of (x_L=0, y_L=0, theta_E=0.5), the Adam\noptimizer runs for 500 steps. Rune differentiates through the entire lens\nequation -- including the division by |theta|, the square root, and the\nmean/variance -- to provide exact gradients that drive convergence to the true\nparameters.\n\n## Try It\n\n1. Increase the noise level from 0.005 to 0.05 and observe how parameter\n uncertainties grow.\n2. Add a shear term (gamma_1, gamma_2) to the lens model for external\n tidal perturbation.\n3. Replace the point-mass with a singular isothermal sphere (SIS) profile\n where the deflection is constant: alpha = theta_E * theta_hat.\n\n## Next Steps\n\nContinue to [10-uncertainty-propagation](../10-uncertainty-propagation/) to\nlearn how to automatically propagate parameter uncertainties through\ncosmological distance calculations using exact AD Jacobians.\n"
},
{
"path": "dev/umbra/examples/09-gravitational-lensing/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega umbra))\n"
},
{
"path": "dev/umbra/examples/09-gravitational-lensing/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Differentiable gravitational lens parameter fitting.\n\n A point-mass gravitational lens produces multiple images of a background\n source. Given the observed image positions, we fit the lens center and\n Einstein radius by requiring that all images map back to the same source\n position. The lens equation and source-plane mapping are expressed as Nx\n tensor operations, making the entire model differentiable through Rune. *)\n\nopen Nx\n\nlet f64 = Nx.float64\n\n(* True lens parameters (for generating synthetic data) *)\nlet true_x_l = 0.1\nlet true_y_l = -0.05\nlet true_theta_e = 1.0\n\n(* Generate synthetic image positions for a quadruply-imaged quasar. Source at\n (0.15, 0.08), lens at (true_x_l, true_y_l). *)\nlet source_x = 0.15\nlet source_y = 0.08\nlet () = Printf.printf \"Differentiable gravitational lens modeling\\n\"\nlet () = Printf.printf \"==========================================\\n\\n\"\n\n(* Solve lens equation: beta = theta - theta_E^2 * theta / |theta|^2 for point\n mass (where theta is relative to lens center). We generate 4 image positions\n by solving analytically + adding noise. *)\nlet img_x, img_y =\n (* For a point mass, images lie along the source-lens axis. Place 4 images at\n realistic positions around the lens. *)\n let dx = source_x -. true_x_l in\n let dy = source_y -. true_y_l in\n let beta = Float.sqrt ((dx *. dx) +. (dy *. dy)) in\n let cos_a = dx /. beta and sin_a = dy /. beta in\n (* Two images along the axis *)\n let theta_p =\n (beta\n +. Float.sqrt ((beta *. beta) +. (4.0 *. true_theta_e *. true_theta_e)))\n /. 2.0\n in\n let theta_m =\n (beta\n -. Float.sqrt ((beta *. beta) +. (4.0 *. true_theta_e *. true_theta_e)))\n /. 2.0\n in\n (* Image positions in 2D (along and perpendicular to axis, with noise) *)\n let noise = 0.005 in\n let x1 = true_x_l +. (theta_p *. cos_a) +. (noise *. 0.3) in\n let y1 = true_y_l +. (theta_p *. sin_a) -. (noise *. 0.2) in\n let x2 = true_x_l +. (theta_m *. cos_a) -. (noise *. 0.5) in\n let y2 = true_y_l +. (theta_m *. sin_a) +. (noise *. 0.4) in\n (* Add two more images from slight perturbation (simulating extended\n source) *)\n let x3 = true_x_l +. (theta_p *. 0.7 *. cos_a) +. (theta_p *. 0.3 *. sin_a) in\n let y3 = true_y_l +. (theta_p *. 0.7 *. sin_a) -. (theta_p *. 0.3 *. cos_a) in\n let x4 = true_x_l -. (theta_p *. 0.5 *. cos_a) +. (theta_p *. 0.4 *. sin_a) in\n let y4 = true_y_l -. (theta_p *. 0.5 *. sin_a) -. (theta_p *. 0.4 *. cos_a) in\n ( create f64 [| 4 |] [| x1; x2; x3; x4 |],\n create f64 [| 4 |] [| y1; y2; y3; y4 |] )\n\n(* Loss: given lens params, map each image back to the source plane. All images\n should map to the same source -> minimize variance of inferred source\n positions. *)\nlet loss params =\n match params with\n | [ x_l; y_l; theta_e ] ->\n (* Displacement from lens center *)\n let dx = sub img_x x_l in\n let dy = sub img_y y_l in\n (* Distance from lens center *)\n let r_sq = add (square dx) (square dy) in\n let r = sqrt r_sq in\n (* Point-mass deflection: alpha = theta_E^2 / r *)\n let alpha = div (square theta_e) r in\n (* Source position for each image: beta = theta - alpha * hat(theta) *)\n let beta_x = sub img_x (mul alpha (div dx r)) in\n let beta_y = sub img_y (mul alpha (div dy r)) in\n (* Variance of source positions (should be ~0 if lens model is correct) *)\n let mean_bx = mean beta_x in\n let mean_by = mean beta_y in\n let var_x = mean (square (sub beta_x mean_bx)) in\n let var_y = mean (square (sub beta_y mean_by)) in\n add var_x var_y\n | _ -> failwith \"expected [x_l; y_l; theta_e]\"\n\nlet () =\n Printf.printf \"True parameters:\\n\";\n Printf.printf \" x_L = %.3f arcsec\\n\" true_x_l;\n Printf.printf \" y_L = %.3f arcsec\\n\" true_y_l;\n Printf.printf \" theta_E = %.3f arcsec\\n\\n\" true_theta_e;\n let algo = Vega.adam (Vega.Schedule.constant 1e-2) in\n let x_l = ref (scalar f64 0.0) in\n let y_l = ref (scalar f64 0.0) in\n let theta_e = ref (scalar f64 0.5) in\n let states =\n [| Vega.init algo !x_l; Vega.init algo !y_l; Vega.init algo !theta_e |]\n in\n let steps = 500 in\n Printf.printf \"%5s %12s %8s %8s %8s\\n\" \"step\" \"loss\" \"x_L\" \"y_L\" \"theta_E\";\n Printf.printf \"%5s %12s %8s %8s %8s\\n\" \"-----\" \"------------\" \"--------\"\n \"--------\" \"--------\";\n let refs = [| x_l; y_l; theta_e |] in\n for i = 0 to steps - 1 do\n let loss_val, grads = Rune.value_and_grads loss [ !x_l; !y_l; !theta_e ] in\n List.iteri\n (fun j g ->\n let p, s = Vega.step states.(j) ~grad:g ~param:!(refs.(j)) in\n refs.(j) := p;\n states.(j) <- s)\n grads;\n if i mod 100 = 0 || i = steps - 1 then\n Printf.printf \"%5d %12.8f %8.4f %8.4f %8.4f\\n\" i (item [] loss_val)\n (item [] !x_l) (item [] !y_l) (item [] !theta_e)\n done;\n Printf.printf \"\\nFitted parameters:\\n\";\n Printf.printf \" x_L = %.4f (true: %.4f)\\n\" (item [] !x_l) true_x_l;\n Printf.printf \" y_L = %.4f (true: %.4f)\\n\" (item [] !y_l) true_y_l;\n Printf.printf \" theta_E = %.4f (true: %.4f)\\n\" (item [] !theta_e)\n true_theta_e\n"
},
{
"path": "dev/umbra/examples/10-uncertainty-propagation/README.md",
"content": "# `10-uncertainty-propagation`\n\nAutomatic uncertainty propagation through cosmological distance calculations.\nPropagates H0 and Omega_m uncertainties through distance modulus using exact\nAD Jacobians via forward-mode differentiation. The linear error propagation\nformula (Sigma_out = J Sigma_in J^T) is validated against Monte Carlo sampling\nwith 50,000 draws.\n\n```bash\ncd dev/umbra\ndune exec --root . examples/10-uncertainty-propagation/main.exe\n```\n\n## What You'll Learn\n\n- Computing exact Jacobians automatically with forward-mode AD (`Rune.jacfwd`)\n- Applying linear error propagation via the Jacobian covariance formula\n- Validating analytical uncertainty estimates with Monte Carlo sampling\n- Propagating scalar uncertainties through cosmological models with JVP\n\n## Key Functions\n\n| Function | Purpose |\n| ---------------------------- | ------------------------------------------------ |\n| `Cosmo.create_flat_lcdm` | Create a flat Lambda-CDM cosmology |\n| `Cosmo.distance_modulus` | Compute distance modulus at a given redshift |\n| `Rune.jacfwd` | Forward-mode Jacobian of a function |\n| `Rune.jvp` | Jacobian-vector product for scalar propagation |\n| `Nx.cholesky` | Cholesky decomposition for MC sampling |\n| `Nx.matmul` | Matrix multiply for J Sigma J^T |\n| `Nx.diag` | Build diagonal covariance from variances |\n\n## How It Works\n\nGiven input parameters with uncertainties (H0 = 70 +/- 1 km/s/Mpc, Omega_m =\n0.30 +/- 0.01), the example propagates these through `Cosmo.distance_modulus`\nat five redshifts (z = 0.1 to 1.0). The propagation uses the standard linear\nformula: Sigma_out = J Sigma_in J^T, where J is the Jacobian of the distance\nmodulus with respect to [H0, Omega_m]. Rather than deriving J analytically,\n`Rune.jacfwd` computes it automatically with just two JVP evaluations (one per\ninput parameter).\n\nFor validation, the example draws 50,000 Monte Carlo samples from the input\ncovariance via Cholesky decomposition, evaluates the model at each sample, and\ncomputes empirical output statistics. Agreement below 1% between AD and MC\nconfirms that linear propagation is accurate for these parameter ranges.\n\nA scalar API demo shows the simpler case: propagating redshift uncertainty\n(z = 0.5 +/- 0.01) through a single `jvp` call, which returns both the output\nvalue and its sensitivity to the input perturbation.\n\n## Try It\n\n1. Add correlation between H0 and Omega_m by putting off-diagonal terms in the\n input covariance matrix.\n2. Increase the uncertainties to see where linear propagation breaks down and\n MC diverges from AD.\n3. Propagate uncertainties through `Cosmo.luminosity_distance` instead of\n distance modulus and compare the relative errors.\n\n## Next Steps\n\nContinue to [11-bayesian-sed](../11-bayesian-sed/) to see how Fisher information\nand Hamiltonian Monte Carlo provide both theoretical bounds and full Bayesian\nposteriors for SED parameter estimation.\n"
},
{
"path": "dev/umbra/examples/10-uncertainty-propagation/dune",
"content": "(executable\n (name main)\n (libraries nx rune umbra))\n"
},
{
"path": "dev/umbra/examples/10-uncertainty-propagation/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Automatic uncertainty propagation through cosmological distances.\n\n Demonstrates propagating H0 and Omega_m uncertainties through\n Umbra.Cosmo.distance_modulus using exact AD Jacobians. The linear error\n propagation formula (Sigma_out = J Sigma_in J^T) is computed automatically\n via Rune.jacfwd. Results are validated against Monte Carlo sampling.\n\n Fisher, propagation, and Monte Carlo are all trivial given Rune's jacfwd --\n no dedicated library needed. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\n(* Redshifts to evaluate *)\nlet redshifts = [| 0.1; 0.3; 0.5; 0.7; 1.0 |]\n\n(* Forward model: given [H0; Omega_m], compute distance modulus at z *)\nlet distance_modulus_at z p =\n let h0 = Nx.reshape [||] (Nx.slice [ I 0 ] p) in\n let om = Nx.reshape [||] (Nx.slice [ I 1 ] p) in\n let cosmo = Cosmo.create_flat_lcdm ~h0 ~omega_m:om in\n Cosmo.distance_modulus ~p:cosmo (Nx.scalar f64 z)\n\n(* Linear error propagation: Sigma_out = J Sigma_in J^T *)\nlet propagate f ~mean ~cov =\n let j = Rune.jacfwd f mean in\n let mean_out = f mean in\n let cov_out = Nx.matmul (Nx.matmul j cov) (Nx.matrix_transpose j) in\n let cov_out = Nx.div_s (Nx.add cov_out (Nx.matrix_transpose cov_out)) 2.0 in\n (mean_out, cov_out)\n\n(* Monte Carlo validation *)\nlet monte_carlo ?(n_samples = 50_000) f ~mean ~cov =\n let n = Nx.numel mean in\n let l = Nx.cholesky cov in\n let z = Nx.randn f64 [| n_samples; n |] in\n let samples = Nx.add (Nx.matmul z (Nx.matrix_transpose l)) mean in\n let y0 = f (Nx.slice [ I 0 ] samples) in\n let m = Nx.numel y0 in\n let outputs = Nx.zeros f64 [| n_samples; m |] in\n Nx.set_slice [ I 0 ] outputs y0;\n for i = 1 to n_samples - 1 do\n Nx.set_slice [ I i ] outputs (f (Nx.slice [ I i ] samples))\n done;\n let mean_out = Nx.mean ~axes:[ 0 ] outputs in\n let centered = Nx.sub outputs mean_out in\n let cov_out =\n Nx.div_s\n (Nx.matmul (Nx.matrix_transpose centered) centered)\n (Float.of_int (n_samples - 1))\n in\n (mean_out, cov_out)\n\nlet () =\n Printf.printf \"Automatic Uncertainty Propagation through Cosmology\\n\";\n Printf.printf \"====================================================\\n\\n\";\n\n (* Parameters with uncertainties *)\n let h0_mean = 70.0 and h0_std = 1.0 in\n let om_mean = 0.30 and om_std = 0.01 in\n\n Printf.printf \"Input parameters:\\n\";\n Printf.printf \" H0 = %.1f +/- %.1f km/s/Mpc\\n\" h0_mean h0_std;\n Printf.printf \" Omega_m = %.2f +/- %.2f\\n\\n\" om_mean om_std;\n\n let mean = Nx.create f64 [| 2 |] [| h0_mean; om_mean |] in\n let std = Nx.create f64 [| 2 |] [| h0_std; om_std |] in\n let cov = Nx.diag (Nx.square std) in\n\n Printf.printf \"%5s %10s %10s %10s %10s\\n\" \"z\" \"mu (AD)\" \"sigma (AD)\"\n \"sigma (MC)\" \"agreement\";\n Printf.printf \"%5s %10s %10s %10s %10s\\n\" \"-----\" \"----------\"\n \"----------\" \"----------\" \"----------\";\n\n Array.iter\n (fun z ->\n (* AD-based propagation *)\n let f p = Nx.reshape [| 1 |] (distance_modulus_at z p) in\n let mean_ad, cov_ad = propagate f ~mean ~cov in\n let mu_ad = item [ 0 ] mean_ad in\n let std_ad = Float.sqrt (item [ 0; 0 ] cov_ad) in\n\n (* Monte Carlo validation *)\n let _, cov_mc = monte_carlo f ~mean ~cov in\n let std_mc = Float.sqrt (item [ 0; 0 ] cov_mc) in\n\n let agreement = Float.abs (std_ad -. std_mc) /. std_mc *. 100.0 in\n Printf.printf \"%5.1f %10.4f %10.4f %10.4f %9.1f%%\\n\" z mu_ad std_ad\n std_mc agreement)\n redshifts;\n\n Printf.printf \"\\n\";\n Printf.printf \"AD uses exact Jacobians (2 JVP calls for 2 parameters).\\n\";\n Printf.printf \"MC uses 50,000 samples for validation.\\n\";\n Printf.printf \"Agreement < 1%% confirms linear propagation is accurate.\\n\";\n\n (* Also demonstrate the simple scalar API *)\n Printf.printf \"\\n--- Scalar API demo ---\\n\\n\";\n Printf.printf \"Propagating z = 0.5 +/- 0.01 through distance_modulus:\\n\";\n let x = Nx.scalar f64 0.5 in\n let y, dy =\n Rune.jvp (fun z -> Cosmo.distance_modulus z) x (Nx.scalar f64 1.0)\n in\n let mu_mean = Nx.item [] y in\n let mu_std = Float.abs (Nx.item [] dy) *. 0.01 in\n Printf.printf \" mu = %.4f +/- %.4f\\n\" mu_mean mu_std\n"
},
{
"path": "dev/umbra/examples/11-bayesian-sed/README.md",
"content": "# `11-bayesian-sed`\n\nFisher information matrix analysis and Hamiltonian Monte Carlo sampling for\nBayesian SED parameter estimation. Computes Cramer-Rao bounds (theoretical\nminimum uncertainties) from the Fisher matrix, then samples the full posterior\nvia HMC through the differentiable spectrum -> extinction -> photometry pipeline.\n\n```bash\ncd dev/umbra\ndune exec --root . examples/11-bayesian-sed/main.exe\n```\n\n## What You'll Learn\n\n- Computing the Fisher information matrix via reverse-mode Jacobians\n- Deriving Cramer-Rao bounds on SED parameters (temperature, extinction)\n- Sampling full Bayesian posteriors with Hamiltonian Monte Carlo\n- Comparing Fisher-predicted vs HMC-sampled uncertainties\n- Building differentiable forward models through tophat bandpasses\n\n## Key Functions\n\n| Function | Purpose |\n| -------------------------- | --------------------------------------------------- |\n| `Rune.jacrev` | Reverse-mode Jacobian for Fisher matrix computation |\n| `Nx.inv` | Matrix inverse for Fisher -> covariance |\n| `Nx.diagonal` | Extract diagonal (marginal variances) |\n| `Spectrum.blackbody` | Generate Planck SED at given temperature |\n| `Extinction.apply` | Apply CCM89 dust reddening |\n| `Photometry.tophat` | Create rectangular bandpass filters |\n| `Photometry.ab_mag` | Compute AB magnitude through a bandpass |\n| `Norn.hmc` | Hamiltonian Monte Carlo posterior sampling |\n\n## How It Works\n\nThe forward model maps two parameters -- log(T) and A_V -- to five broadband\nmagnitudes through the pipeline: `blackbody -> extinction -> ab_mag`. Synthetic\nobservations are generated at T=6500 K, A_V=0.5 with realistic photometric\nerrors (0.03-0.05 mag).\n\nThe Fisher information matrix F = J^T C^-1 J is computed from the Jacobian of\nthe model (via `Rune.jacrev`) and the observational covariance C. Inverting F\ngives the Cramer-Rao lower bound -- the best achievable 1-sigma uncertainty on\neach parameter for a given dataset, regardless of estimation method.\n\nThe example then samples the actual Bayesian posterior using `Norn.hmc`. The\nlog-posterior is a Gaussian likelihood with flat priors, and HMC uses Rune's\ngradients to efficiently explore the parameter space with 500 post-warmup\nsamples. Comparing the HMC posterior width to the Fisher prediction validates\nthat the model is well-behaved: when they agree, the posterior is approximately\nGaussian and the Fisher bound is tight.\n\n## Try It\n\n1. Reduce the photometric errors to 0.01 mag and observe how both Fisher bounds\n and HMC posteriors tighten.\n2. Add a third parameter (redshift) and examine the resulting parameter\n degeneracies in the Fisher matrix.\n3. Replace the flat prior with an informative Gaussian prior on A_V and see\n how the posterior shifts.\n\n## Next Steps\n\nContinue to [12-survey-optimization](../12-survey-optimization/) to see how\ndifferentiable Fisher forecasting enables gradient-based optimization of survey\ndesign parameters for weak gravitational lensing.\n"
},
{
"path": "dev/umbra/examples/11-bayesian-sed/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega norn umbra))\n"
},
{
"path": "dev/umbra/examples/11-bayesian-sed/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Fisher information and HMC sampling for SED parameter estimation.\n\n Demonstrates two capabilities:\n\n 1. Fisher matrix: compute the Cramer-Rao bounds on temperature and extinction\n -- \"how well CAN I constrain these parameters from UGRIZ photometry?\" --\n before taking any data. Computed inline from Rune.jacrev + linear algebra.\n\n 2. HMC sampling: full Bayesian posterior through the differentiable Spectrum\n -> Extinction -> Photometry pipeline, via Norn.hmc. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\n\n(* Bandpasses *)\nlet n_bp = 20\n\nlet bands =\n [\n Photometry.tophat ~lo:(Unit.Length.m 3.0e-7) ~hi:(Unit.Length.m 4.0e-7)\n ~n:n_bp;\n Photometry.tophat ~lo:(Unit.Length.m 4.0e-7) ~hi:(Unit.Length.m 5.5e-7)\n ~n:n_bp;\n Photometry.tophat ~lo:(Unit.Length.m 5.5e-7) ~hi:(Unit.Length.m 7.0e-7)\n ~n:n_bp;\n Photometry.tophat ~lo:(Unit.Length.m 7.0e-7) ~hi:(Unit.Length.m 8.5e-7)\n ~n:n_bp;\n Photometry.tophat ~lo:(Unit.Length.m 8.5e-7) ~hi:(Unit.Length.m 1.0e-6)\n ~n:n_bp;\n ]\n\nlet band_names = [| \"U\"; \"G\"; \"R\"; \"I\"; \"Z\" |]\nlet rv = Nx.scalar f64 3.1\n\n(* Forward model: [log_T, A_V] -> 5 magnitudes *)\nlet model params =\n let log_temp = Nx.reshape [||] (Nx.slice [ I 0 ] params) in\n let av = Nx.reshape [||] (Nx.slice [ I 1 ] params) in\n let temp = Unit.Temperature.of_kelvin (Nx.exp log_temp) in\n let mags =\n List.map\n (fun bp ->\n let wave = Photometry.wavelength bp in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:wave\n |> Extinction.apply (Extinction.ccm89 ~rv) ~av\n |> Spectrum.as_flux_density\n in\n Photometry.ab_mag bp sed)\n bands\n in\n Nx.stack ~axis:0 mags\n\n(* True parameters *)\nlet true_log_temp = Float.log 6500.0\nlet true_av = 0.5\nlet true_params = Nx.create f64 [| 2 |] [| true_log_temp; true_av |]\n\n(* Synthetic observations *)\nlet obs_errs = Nx.create f64 [| 5 |] [| 0.05; 0.03; 0.03; 0.04; 0.05 |]\n\nlet obs_mags =\n let true_mags = model true_params in\n let noise = Nx.create f64 [| 5 |] [| 0.03; -0.02; 0.01; -0.01; 0.02 |] in\n Nx.add true_mags noise\n\n(* Fisher information: F = J^T C^-1 J *)\nlet fisher f ~params ~obs_cov =\n let j = Rune.jacrev f params in\n let jt = Nx.matrix_transpose j in\n Nx.matmul (Nx.matmul jt (Nx.inv obs_cov)) j\n\n(* Cramer-Rao bounds: sigma = sqrt(diag(F^-1)) *)\nlet marginal_sigma f = Nx.sqrt (Nx.diagonal (Nx.inv f))\n\nlet () =\n Printf.printf \"Fisher Information & HMC for SED Fitting\\n\";\n Printf.printf \"=========================================\\n\\n\";\n\n Printf.printf \"True parameters:\\n\";\n Printf.printf \" T = %.0f K (log_T = %.4f)\\n\" (Float.exp true_log_temp)\n true_log_temp;\n Printf.printf \" A_V = %.2f\\n\\n\" true_av;\n\n Printf.printf \"Observed magnitudes:\\n\";\n Array.iteri\n (fun i name ->\n Printf.printf \" %s = %.3f +/- %.3f\\n\" name (item [ i ] obs_mags)\n (item [ i ] obs_errs))\n band_names;\n Printf.printf \"\\n\";\n\n (* --- Fisher Information --- *)\n Printf.printf \"=== Fisher Information ===\\n\\n\";\n\n let obs_cov = Nx.diag (Nx.square obs_errs) in\n let f = fisher model ~params:true_params ~obs_cov in\n let sigma = marginal_sigma f in\n\n Printf.printf \"Fisher matrix:\\n\";\n Printf.printf \" F = [[ %10.2f %10.2f ]\\n\"\n (item [ 0; 0 ] f)\n (item [ 0; 1 ] f);\n Printf.printf \" [ %10.2f %10.2f ]]\\n\\n\"\n (item [ 1; 0 ] f)\n (item [ 1; 1 ] f);\n\n Printf.printf \"Cramer-Rao bounds (best achievable 1-sigma):\\n\";\n let sigma_log_t = item [ 0 ] sigma in\n let sigma_av = item [ 1 ] sigma in\n Printf.printf \" sigma(log_T) = %.4f -> sigma(T) ~ %.0f K\\n\" sigma_log_t\n (sigma_log_t *. Float.exp true_log_temp);\n Printf.printf \" sigma(A_V) = %.4f\\n\\n\" sigma_av;\n\n (* --- HMC Sampling --- *)\n Printf.printf \"=== HMC Posterior Sampling ===\\n\\n\";\n\n (* Log-posterior: Gaussian likelihood, flat prior *)\n let log_posterior params =\n let pred = model params in\n let residuals = Nx.div (Nx.sub pred obs_mags) obs_errs in\n Nx.mul_s (Nx.sum (Nx.square residuals)) (-0.5)\n in\n\n let init = Nx.create f64 [| 2 |] [| Float.log 7000.0; 0.3 |] in\n let result =\n Norn.hmc ~step_size:0.001 ~num_leapfrog:10 ~num_warmup:200 ~n:500\n log_posterior init\n in\n\n Printf.printf \"HMC diagnostics:\\n\";\n Printf.printf \" Accept rate: %.1f%%\\n\\n\" (result.stats.accept_rate *. 100.);\n\n (* Sample statistics *)\n let sample_mean = Nx.mean ~axes:[ 0 ] result.samples in\n let centered = Nx.sub result.samples sample_mean in\n let sample_cov =\n Nx.div_s\n (Nx.matmul (Nx.matrix_transpose centered) centered)\n (Float.of_int 499)\n in\n let sample_std = Nx.sqrt (Nx.diagonal sample_cov) in\n\n let hmc_log_t = item [ 0 ] sample_mean in\n let hmc_av = item [ 1 ] sample_mean in\n let hmc_sigma_log_t = item [ 0 ] sample_std in\n let hmc_sigma_av = item [ 1 ] sample_std in\n\n Printf.printf \"Posterior (HMC):\\n\";\n Printf.printf \" log_T = %.4f +/- %.4f -> T ~ %.0f K\\n\" hmc_log_t\n hmc_sigma_log_t (Float.exp hmc_log_t);\n Printf.printf \" A_V = %.4f +/- %.4f\\n\\n\" hmc_av hmc_sigma_av;\n\n (* --- Comparison --- *)\n Printf.printf \"=== Fisher vs HMC Comparison ===\\n\\n\";\n Printf.printf \" %12s %10s %10s\\n\" \"\" \"Fisher s\" \"HMC s\";\n Printf.printf \" %12s %10s %10s\\n\" \"------------\" \"----------\" \"----------\";\n Printf.printf \" %12s %10.4f %10.4f\\n\" \"s(log_T)\" sigma_log_t\n hmc_sigma_log_t;\n Printf.printf \" %12s %10.4f %10.4f\\n\\n\" \"s(A_V)\" sigma_av hmc_sigma_av;\n\n Printf.printf \"Fisher gives the theoretical minimum uncertainty.\\n\";\n Printf.printf \"HMC gives the actual posterior width.\\n\";\n Printf.printf \"Agreement confirms the model is well-behaved (near-linear).\\n\"\n"
},
{
"path": "dev/umbra/examples/12-survey-optimization/README.md",
"content": "# `12-survey-optimization`\n\nDifferentiable survey optimization for a Stage IV weak lensing survey. Uses\nexact autodiff gradients to optimize survey parameters that minimize the\nuncertainty on S8 = sigma8 * sqrt(Omega_m / 0.3), replacing traditional grid\nsearch with gradient-based Fisher forecasting. Demonstrates both a single-bin\narea/depth tradeoff and joint optimization of sky fraction with tomographic bin\nedges.\n\n```bash\ncd dev/umbra\ndune exec --root . examples/12-survey-optimization/main.exe\n```\n\n## What You'll Learn\n\n- Computing differentiable Fisher information matrices for survey forecasting\n- Optimizing the area/depth tradeoff for sky coverage vs galaxy density\n- Jointly optimizing sky fraction and tomographic bin edges with gradient descent\n- Using sigmoid-windowed bins for smooth gradient flow through discrete boundaries\n- Comparing gradient-based optimization against brute-force grid search\n\n## Key Functions\n\n| Function | Purpose |\n| --------------------------- | ---------------------------------------------------- |\n| `Survey.angular_cl` | Compute angular power spectra for tracer pairs |\n| `Survey.weak_lensing` | Create a weak lensing tracer from n(z) |\n| `Survey.smail` | Smail redshift distribution for source galaxies |\n| `Cosmo.planck18` | Planck 2018 fiducial cosmology |\n| `Cosmo.linear_power` | Linear matter power spectrum P(k, z) |\n| `Cosmo.comoving_distance` | Comoving distance for lensing kernel computation |\n| `Rune.value_and_grad` | Loss and gradient for survey parameter optimization |\n| `Vega.adam` | Adam optimizer for continuous parameter search |\n\n## How It Works\n\nPart 1 tackles the area/depth tradeoff for a single tomographic bin. A fixed\ngalaxy budget (n_gal * f_sky = constant) means wider surveys are shallower. The\nFisher matrix for [Omega_m, sigma8] is computed from Limber-integrated angular\npower spectra, with shape noise that depends on galaxy density. The objective\nfunction -- sigma(S8) derived from the 2x2 Fisher inverse -- is fully\ndifferentiable through f_sky via sigmoid parameterization. Adam finds the\noptimal sky fraction in 300 steps with exact gradients, verified by a\nfinite-difference check.\n\nPart 2 extends to joint optimization of sky fraction and two tomographic bin\nedges that divide galaxies into three redshift bins. The bin boundaries use\nsigmoid window functions (with width delta=0.03) so that gradients flow smoothly\nthrough the discrete bin assignment. Narrower bins concentrate signal but\nincrease shot noise; the optimizer balances this tradeoff automatically. The\nLimber integral uses precomputed cosmological grids (comoving distances, Hubble\nrates, power spectra) evaluated at five cosmology perturbations for numerical\nderivatives of C_l with respect to Omega_m and sigma8, while gradients with\nrespect to survey parameters (f_sky, z1, z2) flow through Rune's autodiff.\n\nA brute-force grid search over 12 x 15 x 15 = 2700 parameter combinations\nvalidates the gradient result, demonstrating that 500 Adam steps achieve equal\nor better precision with orders of magnitude fewer function evaluations.\n\n## Try It\n\n1. Increase the galaxy budget from 10 to 50 gal/arcmin2 and observe how the\n optimal sky fraction shifts toward wider coverage.\n2. Add a fourth tomographic bin and compare the improvement in sigma(S8).\n3. Replace the Smail n(z) with a sharper distribution and see how the optimal\n bin edges respond.\n\n## Next Steps\n\nThis is the final example in the Umbra series. For earlier topics, revisit\n[01-constants-and-units](../01-constants-and-units/) for physical constants and\nunit handling, or [05-sed-fitting](../05-sed-fitting/) for the foundations of\ndifferentiable spectral energy distribution fitting that this example builds on.\n"
},
{
"path": "dev/umbra/examples/12-survey-optimization/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega umbra))\n"
},
{
"path": "dev/umbra/examples/12-survey-optimization/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Differentiable survey optimization via autodiff gradients through the Fisher\n information matrix.\n\n Traditional survey optimization uses grid search over discrete Fisher\n forecasts. Umbra's fully differentiable cosmology pipeline enables\n gradient-based continuous optimization: compute Fisher(survey_params) and\n minimize sigma(S8) with respect to survey parameters using exact autodiff\n gradients from Rune.\n\n Part 1: Area/depth tradeoff -- optimize f_sky with fixed n(z) shape. Part 2:\n Joint area + bin edge optimization -- optimize f_sky and tomographic bin\n edges simultaneously, with gradients flowing through the lensing kernel\n computation via differentiable n(z) windowing. *)\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\nlet sigma_e = 0.26\nlet steradian_to_arcmin2 = 11818102.86004228\nlet c_km_s = 299792.458\nlet h0_ref = 100.0\n\n(* Fiducial cosmology *)\nlet p_fid = Cosmo.planck18\nlet omega_m_fid = Nx.item [] (Cosmo.omega_m p_fid)\nlet sigma8_fid = Nx.item [] (Cosmo.sigma8 p_fid)\n\n(* S8 = sigma8 * sqrt(omega_m / 0.3) -- derivatives at fiducial *)\nlet ds8_dom = sigma8_fid /. (2.0 *. Float.sqrt (0.3 *. omega_m_fid))\nlet ds8_ds8 = Float.sqrt (omega_m_fid /. 0.3)\n\n(* ell weights: (2*ell+1) * dell / 2 *)\nlet ell_weights ell =\n let n_ell = (Nx.shape ell).(0) in\n let dell =\n Array.init n_ell (fun l ->\n if l = 0 then Nx.item [ 1 ] ell -. Nx.item [ 0 ] ell\n else if l = n_ell - 1 then Nx.item [ l ] ell -. Nx.item [ l - 1 ] ell\n else 0.5 *. (Nx.item [ l + 1 ] ell -. Nx.item [ l - 1 ] ell))\n in\n Nx.create f64 [| n_ell |]\n (Array.init n_ell (fun l ->\n ((2.0 *. Nx.item [ l ] ell) +. 1.0) *. dell.(l) /. 2.0))\n\n(* Compute dCl/d(theta) via central finite differences *)\nlet finite_diff_cl ~ell ~tracers ~param_name ~set_param ~fid_val ~eps =\n let p_plus = set_param (scalar f64 (fid_val +. eps)) p_fid in\n let p_minus = set_param (scalar f64 (fid_val -. eps)) p_fid in\n let cl_p =\n Survey.Cls.to_tensor\n (Survey.angular_cl ~p:p_plus ~power:Survey.linear ~ell tracers)\n in\n let cl_m =\n Survey.Cls.to_tensor\n (Survey.angular_cl ~p:p_minus ~power:Survey.linear ~ell tracers)\n in\n let dcl = Nx.div_s (Nx.sub cl_p cl_m) (2.0 *. eps) in\n Printf.printf \" dCl/d(%-8s): max=%.3e\\n\" param_name (Nx.item [] (Nx.max dcl));\n dcl\n\n(* 2x2 analytical Fisher inverse -> sigma(S8) -- all differentiable *)\nlet sigma_s8_from_fisher f11 f12 f22 =\n let det = Nx.sub (Nx.mul f11 f22) (Nx.mul f12 f12) in\n let a = scalar f64 ds8_dom and b = scalar f64 ds8_ds8 in\n let sigma_sq =\n Nx.div\n (Nx.add\n (Nx.sub\n (Nx.mul f22 (Nx.mul a a))\n (Nx.mul_s (Nx.mul f12 (Nx.mul a b)) 2.0))\n (Nx.mul f11 (Nx.mul b b)))\n det\n in\n Nx.sqrt sigma_sq\n\n(* ===================================================================== *)\n(* Part 1: Area/depth tradeoff (single bin) *)\n(* ===================================================================== *)\n\nlet part1 () =\n Printf.printf \"--- Part 1: Area/Depth Tradeoff (1 bin) ---\\n\\n\";\n let budget = 10.0 in\n let ell = Nx.logspace f64 1.0 3.0 30 in\n let w_ell = ell_weights ell in\n\n let nz = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.3 () in\n let wl = Survey.weak_lensing nz in\n\n Printf.printf \"Precomputing signal derivatives...\\n\";\n let cl_fid =\n Survey.Cls.to_tensor\n (Survey.angular_cl ~p:p_fid ~power:Survey.linear ~ell [ wl ])\n in\n let cl_fid_flat = Nx.flatten cl_fid in\n let eps = 1e-4 in\n let dcl_dom =\n Nx.flatten\n (finite_diff_cl ~ell ~tracers:[ wl ] ~param_name:\"omega_m\"\n ~set_param:(fun v p -> Cosmo.set_t ~omega_m:v p)\n ~fid_val:omega_m_fid ~eps)\n in\n let dcl_ds8 =\n Nx.flatten\n (finite_diff_cl ~ell ~tracers:[ wl ] ~param_name:\"sigma8\"\n ~set_param:(fun v p -> Cosmo.set_t ~sigma8:v p)\n ~fid_val:sigma8_fid ~eps)\n in\n Printf.printf \"\\n\";\n\n let objective log_f_sky =\n let f_sky = Nx.sigmoid log_f_sky in\n let n_gal = Nx.div (scalar f64 budget) f_sky in\n let noise =\n Nx.div\n (scalar f64 (sigma_e *. sigma_e))\n (Nx.mul_s n_gal steradian_to_arcmin2)\n in\n let cl_obs = Nx.add cl_fid_flat noise in\n let cl_obs_sq = Nx.mul cl_obs cl_obs in\n let weighted_dom =\n Nx.div (Nx.mul w_ell (Nx.mul dcl_dom dcl_dom)) cl_obs_sq\n in\n let weighted_ds8 =\n Nx.div (Nx.mul w_ell (Nx.mul dcl_ds8 dcl_ds8)) cl_obs_sq\n in\n let weighted_x = Nx.div (Nx.mul w_ell (Nx.mul dcl_dom dcl_ds8)) cl_obs_sq in\n let f11 = Nx.mul f_sky (Nx.sum weighted_dom) in\n let f12 = Nx.mul f_sky (Nx.sum weighted_x) in\n let f22 = Nx.mul f_sky (Nx.sum weighted_ds8) in\n sigma_s8_from_fisher f11 f12 f22\n in\n\n (* Gradient check *)\n let log_f_sky_init = scalar f64 0.0 in\n let v0, g0 = Rune.value_and_grad objective log_f_sky_init in\n let fd_eps = 1e-5 in\n let vp = item [] (objective (scalar f64 fd_eps)) in\n let vm = item [] (objective (scalar f64 (-.fd_eps))) in\n let fd = (vp -. vm) /. (2.0 *. fd_eps) in\n Printf.printf \"Gradient check: AD=%.6e FD=%.6e rel=%.2e\\n\\n\" (item [] g0) fd\n (Float.abs (item [] g0 -. fd) /. Float.abs fd);\n\n let f_sky_0 = 1.0 /. (1.0 +. Float.exp (-0.0)) in\n Printf.printf \"Initial: f_sky=%.3f n_gal=%.1f sigma(S8)=%.6f\\n\" f_sky_0\n (budget /. f_sky_0) (item [] v0);\n\n let algo = Vega.adam (Vega.Schedule.constant 0.01) in\n let log_f_sky = ref log_f_sky_init in\n let state = ref (Vega.init algo !log_f_sky) in\n let best_sigma = ref (item [] v0) in\n let best_f_sky = ref f_sky_0 in\n Printf.printf \"\\n%5s %8s %8s %10s\\n\" \"step\" \"f_sky\" \"n_gal\" \"sigma(S8)\";\n Printf.printf \"%5s %8s %8s %10s\\n\" \"-----\" \"--------\" \"--------\"\n \"----------\";\n let steps = 300 in\n for i = 0 to steps - 1 do\n let sigma_val, grad = Rune.value_and_grad objective !log_f_sky in\n let p, s = Vega.step !state ~grad ~param:!log_f_sky in\n log_f_sky := p;\n state := s;\n let f_sky_cur = 1.0 /. (1.0 +. Float.exp (-.item [] !log_f_sky)) in\n let sigma_cur = item [] sigma_val in\n if sigma_cur < !best_sigma then begin\n best_sigma := sigma_cur;\n best_f_sky := f_sky_cur\n end;\n if i mod 50 = 0 || i = steps - 1 then\n Printf.printf \"%5d %8.4f %8.1f %10.6f\\n\" i f_sky_cur\n (budget /. f_sky_cur) sigma_cur\n done;\n Printf.printf \"\\nOptimal: f_sky=%.4f n_gal=%.1f gal/arcmin2\\n\" !best_f_sky\n (budget /. !best_f_sky);\n Printf.printf \"Improvement: sigma(S8) reduced by %.1f%% vs initial\\n\\n\"\n ((1.0 -. (!best_sigma /. item [] v0)) *. 100.0)\n\n(* ===================================================================== *)\n(* Part 2: Joint area + bin edge optimization (3 bins) *)\n(* ===================================================================== *)\n\n(* Precomputed cosmological grids -- expensive, done once per cosmology. *)\ntype cosmo_grid = {\n n_z : int;\n dz : float;\n z_arr : float array;\n z_vec : Nx.float64_t;\n chi_safe : Nx.float64_t;\n omega_m_t : Nx.float64_t;\n integ_weight : Nx.float64_t;\n w_pk : Nx.float64_t;\n ell_factor_sq : Nx.float64_t;\n}\n\nlet precompute_grid ~p ~ell =\n let zmax = 3.0 in\n let n_z = 50 in\n let dz = zmax /. Float.of_int (n_z - 1) in\n let z_arr = Array.init n_z (fun i -> Float.of_int i *. dz) in\n z_arr.(0) <- 1e-6;\n let z_vec = Nx.create f64 [| n_z |] z_arr in\n let sw =\n Array.init n_z (fun i ->\n if i = 0 || i = n_z - 1 then 1.0 else if i mod 2 = 1 then 4.0 else 2.0)\n in\n let simpson_w = Nx.mul_s (Nx.create f64 [| n_z |] sw) (dz /. 3.0) in\n let h_t = Nx.item [] (Nx.div (Cosmo.h0 p) (Nx.scalar f64 h0_ref)) in\n let chi_vec =\n Nx.create f64 [| n_z |]\n (Array.init n_z (fun j ->\n let z_t = Nx.scalar f64 z_arr.(j) in\n let chi =\n Nx.item [] (Unit.Length.in_mpc (Cosmo.comoving_distance ~p z_t))\n in\n chi *. h_t))\n in\n let chi_safe = Nx.clamp ~min:1e-10 chi_vec in\n let h_vec_f =\n Array.init n_z (fun j ->\n Nx.item [] (Cosmo.hubble ~p (Nx.scalar f64 z_arr.(j))))\n in\n let dchi_dz_vec =\n Nx.create f64 [| n_z |]\n (Array.init n_z (fun j -> h_t *. c_km_s /. h_vec_f.(j)))\n in\n let omega_m_t = Nx.scalar f64 (Nx.item [] (Cosmo.omega_m p)) in\n let integ_weight =\n Nx.create f64 [| n_z |]\n (Array.init n_z (fun j ->\n let sw_j = Nx.item [ j ] simpson_w in\n let dchi_j = Nx.item [ j ] dchi_dz_vec in\n let chi_j = Nx.item [ j ] chi_safe in\n sw_j *. dchi_j /. (chi_j *. chi_j) /. (c_km_s *. c_km_s)))\n in\n let pk_grid =\n Nx.stack\n (List.init n_z (fun j ->\n let z_t = Nx.scalar f64 z_arr.(j) in\n let chi_j = Nx.item [ j ] chi_safe in\n let k_vec = Nx.div_s (Nx.add_s ell 0.5) chi_j in\n Cosmo.linear_power ~p k_vec z_t))\n in\n let w_pk =\n Nx.mul\n (Nx.reshape [| n_z; 1 |]\n (Nx.create f64 [| n_z |]\n (Array.init n_z (fun j -> Nx.item [ j ] integ_weight))))\n pk_grid\n in\n let l = ell in\n let num =\n Nx.mul\n (Nx.mul (Nx.sub_s l 1.0) l)\n (Nx.mul (Nx.add_s l 1.0) (Nx.add_s l 2.0))\n in\n let den = Nx.mul (Nx.add_s l 0.5) (Nx.add_s l 0.5) in\n let ell_factor = Nx.div (Nx.sqrt (Nx.abs num)) den in\n let ell_factor_sq = Nx.mul ell_factor ell_factor in\n {\n n_z;\n dz;\n z_arr;\n z_vec;\n chi_safe;\n omega_m_t;\n integ_weight;\n w_pk;\n ell_factor_sq;\n }\n\n(* Reverse cumulative trapezoidal sum *)\nlet rev_cumtrapz f_vec n dz =\n let left = Nx.slice [ R (0, n - 1) ] f_vec in\n let right = Nx.slice [ R (1, n) ] f_vec in\n let mid = Nx.mul_s (Nx.add left right) (0.5 *. dz) in\n let partial = Nx.flip (Nx.cumsum ~axis:0 (Nx.flip mid)) in\n Nx.concatenate [ partial; Nx.zeros f64 [| 1 |] ]\n\n(* Fast WL-only angular Cl from precomputed cosmo grid + pre-evaluated n(z)\n tensors. nz_tensors are [n_z] tensors, one per bin, evaluated on the z grid.\n Differentiable through the n(z) values. *)\nlet fast_wl_cl grid nz_tensors =\n let n_z = grid.n_z and dz = grid.dz in\n let n_bins = Array.length nz_tensors in\n (* Build WL kernels *)\n let prefactor =\n Nx.mul_s grid.omega_m_t (3.0 *. h0_ref *. h0_ref /. (2.0 *. c_km_s))\n in\n let one_plus_z = Nx.add_s grid.z_vec 1.0 in\n let kernels =\n Array.init n_bins (fun b ->\n let nz_t = nz_tensors.(b) in\n let a_vec = rev_cumtrapz nz_t n_z dz in\n let nz_over_chi = Nx.div nz_t grid.chi_safe in\n let b_vec = rev_cumtrapz nz_over_chi n_z dz in\n let g_vec = Nx.sub a_vec (Nx.mul grid.chi_safe b_vec) in\n Nx.mul prefactor (Nx.mul one_plus_z (Nx.mul grid.chi_safe g_vec)))\n in\n (* Limber integration for all pairs *)\n let pairs = ref [] in\n for i = 0 to n_bins - 1 do\n for j = i to n_bins - 1 do\n pairs := (i, j) :: !pairs\n done\n done;\n let pairs = List.rev !pairs in\n Nx.stack\n (List.map\n (fun (i, j) ->\n let ki = Nx.reshape [| n_z; 1 |] kernels.(i) in\n let kj = Nx.reshape [| n_z; 1 |] kernels.(j) in\n let integrand = Nx.mul (Nx.mul ki kj) grid.w_pk in\n Nx.mul grid.ell_factor_sq (Nx.sum ~axes:[ 0 ] integrand))\n pairs)\n\n(* Parent n(z): Smail distribution, evaluated as float *)\nlet parent_nz =\n let a = 2.0 and b = 1.5 and z0 = 0.3 in\n let raw z_f = (z_f ** a) *. Float.exp (-.((z_f /. z0) ** b)) in\n let norm =\n let n = 256 in\n let h = 3.0 /. Float.of_int n in\n let s = ref (raw 1e-6 +. raw 3.0) in\n for i = 1 to n - 1 do\n let x = Float.of_int i *. h in\n let w = if i mod 2 = 1 then 4.0 else 2.0 in\n s := !s +. (w *. raw x)\n done;\n !s *. h /. 3.0\n in\n fun z_f -> raw z_f /. norm\n\n(* Build a differentiable bin n(z) with smooth sigmoid edges *)\nlet make_bin_eval z_lo z_hi delta z =\n let parent_val = parent_nz (Nx.item [] z) in\n if parent_val < 1e-30 then scalar f64 0.0\n else\n let lo_gate = Nx.sigmoid (Nx.div_s (Nx.sub z z_lo) delta) in\n let hi_gate = Nx.sigmoid (Nx.div_s (Nx.sub z_hi z) delta) in\n Nx.mul_s (Nx.mul lo_gate hi_gate) parent_val\n\nlet part2 () =\n Printf.printf \"--- Part 2: Joint Area + Bin Edges (3 bins) ---\\n\\n\";\n let budget = 10.0 in\n let ell = Nx.logspace f64 1.0 3.0 20 in\n let w_ell = ell_weights ell in\n let eps = 1e-4 in\n let delta = 0.03 in\n\n Printf.printf \"Precomputing cosmo grids (fiducial + 4 perturbations)...\\n\";\n let grid_fid = precompute_grid ~p:p_fid ~ell in\n let grid_p_om =\n precompute_grid\n ~p:(Cosmo.set_t ~omega_m:(scalar f64 (omega_m_fid +. eps)) p_fid)\n ~ell\n in\n let grid_m_om =\n precompute_grid\n ~p:(Cosmo.set_t ~omega_m:(scalar f64 (omega_m_fid -. eps)) p_fid)\n ~ell\n in\n let grid_p_s8 =\n precompute_grid\n ~p:(Cosmo.set_t ~sigma8:(scalar f64 (sigma8_fid +. eps)) p_fid)\n ~ell\n in\n let grid_m_s8 =\n precompute_grid\n ~p:(Cosmo.set_t ~sigma8:(scalar f64 (sigma8_fid -. eps)) p_fid)\n ~ell\n in\n Printf.printf \"Done.\\n\\n\";\n\n let n_z = grid_fid.n_z in\n let z_arr = grid_fid.z_arr in\n let dz = grid_fid.dz in\n\n let objective params =\n let log_f_sky = Nx.get [ 0 ] params in\n let z1 = Nx.get [ 1 ] params in\n let z2 = Nx.get [ 2 ] params in\n let f_sky = Nx.sigmoid log_f_sky in\n let n_gal = Nx.div (scalar f64 budget) f_sky in\n\n (* Differentiable n(z) bin functions *)\n let nz_funs =\n [|\n make_bin_eval (scalar f64 0.0) z1 delta;\n make_bin_eval z1 z2 delta;\n make_bin_eval z2 (scalar f64 3.0) delta;\n |]\n in\n\n (* Evaluate n(z) on z grid -- differentiable through bin edges *)\n let nz_tensors =\n Array.init 3 (fun b ->\n Nx.stack\n (List.init n_z (fun j -> nz_funs.(b) (Nx.scalar f64 z_arr.(j)))))\n in\n\n (* Galaxy fraction per bin: integral of window_i(z) n(z) dz. Parent n(z) is\n normalized so this gives the fraction of total galaxies in each bin.\n Differentiable through bin edges -- narrow bins get fewer galaxies. *)\n let gal_fracs =\n Array.init 3 (fun b ->\n let nz_t = nz_tensors.(b) in\n let left = Nx.slice [ R (0, n_z - 2) ] nz_t in\n let right = Nx.slice [ R (1, n_z - 1) ] nz_t in\n Nx.mul_s (Nx.sum (Nx.add left right)) (0.5 *. dz))\n in\n\n (* Per-bin noise: sigma_e^2 / (n_gal_bin * ster) where n_gal_bin = n_gal *\n f_i. Bins with fewer galaxies have higher shot noise. *)\n let noise_per_bin =\n Array.init 3 (fun b ->\n Nx.div\n (scalar f64 (sigma_e *. sigma_e))\n (Nx.mul_s (Nx.mul n_gal gal_fracs.(b)) steradian_to_arcmin2))\n in\n\n (* Fast Cl from precomputed grids -- only n(z) -> kernel is traced *)\n let cl_fid = fast_wl_cl grid_fid nz_tensors in\n let cl_p_om = fast_wl_cl grid_p_om nz_tensors in\n let cl_m_om = fast_wl_cl grid_m_om nz_tensors in\n let cl_p_s8 = fast_wl_cl grid_p_s8 nz_tensors in\n let cl_m_s8 = fast_wl_cl grid_m_s8 nz_tensors in\n let dcl_dom = Nx.div_s (Nx.sub cl_p_om cl_m_om) (2.0 *. eps) in\n let dcl_ds8 = Nx.div_s (Nx.sub cl_p_s8 cl_m_s8) (2.0 *. eps) in\n\n (* Full Fisher via Tr[C^-1 dC/dtheta_i C^-1 dC/dtheta_j] with analytical 3x3\n inverse. Vectorized over ell: each matrix element is a [n_ell] tensor. *)\n let n_bins = 3 in\n\n (* Pair index: (i,j) -> spectrum row in cl arrays. Ordering: (0,0)=0,\n (0,1)=1, (0,2)=2, (1,1)=3, (1,2)=4, (2,2)=5 *)\n let pidx i j =\n let a, b = if i <= j then (i, j) else (j, i) in\n (a * ((2 * n_bins) - a - 1) / 2) + b\n in\n\n (* Build 3x3 C(ell) = Cl + N, stored as flat [9] of [n_ell] tensors *)\n let c =\n Array.init 9 (fun idx ->\n let i = idx / 3 and j = idx mod 3 in\n let cl_ij = Nx.slice [ I (pidx i j) ] cl_fid in\n if i = j then Nx.add cl_ij noise_per_bin.(i) else cl_ij)\n in\n\n (* 3x3 inverse via cofactors / determinant *)\n let det =\n Nx.add\n (Nx.sub\n (Nx.mul c.(0) (Nx.sub (Nx.mul c.(4) c.(8)) (Nx.mul c.(5) c.(7))))\n (Nx.mul c.(1) (Nx.sub (Nx.mul c.(3) c.(8)) (Nx.mul c.(5) c.(6)))))\n (Nx.mul c.(2) (Nx.sub (Nx.mul c.(3) c.(7)) (Nx.mul c.(4) c.(6))))\n in\n let ci = Array.make 9 (scalar f64 0.0) in\n ci.(0) <- Nx.div (Nx.sub (Nx.mul c.(4) c.(8)) (Nx.mul c.(5) c.(7))) det;\n ci.(1) <- Nx.div (Nx.sub (Nx.mul c.(2) c.(7)) (Nx.mul c.(1) c.(8))) det;\n ci.(2) <- Nx.div (Nx.sub (Nx.mul c.(1) c.(5)) (Nx.mul c.(2) c.(4))) det;\n ci.(3) <- Nx.div (Nx.sub (Nx.mul c.(5) c.(6)) (Nx.mul c.(3) c.(8))) det;\n ci.(4) <- Nx.div (Nx.sub (Nx.mul c.(0) c.(8)) (Nx.mul c.(2) c.(6))) det;\n ci.(5) <- Nx.div (Nx.sub (Nx.mul c.(2) c.(3)) (Nx.mul c.(0) c.(5))) det;\n ci.(6) <- Nx.div (Nx.sub (Nx.mul c.(3) c.(7)) (Nx.mul c.(4) c.(6))) det;\n ci.(7) <- Nx.div (Nx.sub (Nx.mul c.(1) c.(6)) (Nx.mul c.(0) c.(7))) det;\n ci.(8) <- Nx.div (Nx.sub (Nx.mul c.(0) c.(4)) (Nx.mul c.(1) c.(3))) det;\n\n (* Build dC/dtheta matrices: symmetric, no noise term *)\n let dc_om =\n Array.init 9 (fun idx ->\n Nx.slice [ I (pidx (idx / 3) (idx mod 3)) ] dcl_dom)\n in\n let dc_s8 =\n Array.init 9 (fun idx ->\n Nx.slice [ I (pidx (idx / 3) (idx mod 3)) ] dcl_ds8)\n in\n\n (* 3x3 matmul: (AB)_ij = sum_k A_ik B_kj, vectorized over ell *)\n let mm a b =\n Array.init 9 (fun idx ->\n let i = idx / 3 and j = idx mod 3 in\n Nx.add\n (Nx.add (Nx.mul a.(i * 3) b.(j)) (Nx.mul a.((i * 3) + 1) b.(3 + j)))\n (Nx.mul a.((i * 3) + 2) b.(6 + j)))\n in\n\n (* Tr[AB] = sum_ij A_ij B_ji, returns [n_ell] tensor *)\n let tr a b =\n let t = ref (Nx.mul a.(0) b.(0)) in\n for i = 0 to 2 do\n for j = 0 to 2 do\n if i > 0 || j > 0 then\n t := Nx.add !t (Nx.mul a.((i * 3) + j) b.((j * 3) + i))\n done\n done;\n !t\n in\n\n (* D1 = C^-1 dC/d(Omega_m), D2 = C^-1 dC/d(sigma8) *)\n let d1 = mm ci dc_om in\n let d2 = mm ci dc_s8 in\n\n (* F_ij = f_sky * sum_ell w_ell * Tr[D_i D_j] *)\n let f11 = Nx.mul f_sky (Nx.sum (Nx.mul w_ell (tr d1 d1))) in\n let f12 = Nx.mul f_sky (Nx.sum (Nx.mul w_ell (tr d1 d2))) in\n let f22 = Nx.mul f_sky (Nx.sum (Nx.mul w_ell (tr d2 d2))) in\n sigma_s8_from_fisher f11 f12 f22\n in\n\n let params = Nx.create f64 [| 3 |] [| -1.1; 0.5; 1.0 |] in\n Printf.printf \"Computing initial sigma(S8)...\\n\";\n let v0 = item [] (objective params) in\n let f_sky_0 = 1.0 /. (1.0 +. Float.exp 1.1) in\n Printf.printf\n \"Initial: f_sky=%.3f bins=[0, 0.50, 1.00, 3.0] sigma(S8)=%.6f\\n\\n\" f_sky_0\n v0;\n\n let algo = Vega.adam (Vega.Schedule.constant 0.03) in\n let params = ref params in\n let state = ref (Vega.init algo !params) in\n let best_sigma = ref v0 in\n let best_params = ref !params in\n Printf.printf \"%5s %8s %8s %8s %10s\\n\" \"step\" \"f_sky\" \"z1\" \"z2\"\n \"sigma(S8)\";\n Printf.printf \"%5s %8s %8s %8s %10s\\n\" \"-----\" \"--------\" \"--------\"\n \"--------\" \"----------\";\n let steps = 500 in\n for i = 0 to steps - 1 do\n let sigma_val, grad = Rune.value_and_grad objective !params in\n let p, s = Vega.step !state ~grad ~param:!params in\n let z1 = Float.max 0.1 (Float.min 2.8 (item [ 1 ] p)) in\n let z2 = Float.max (z1 +. 0.1) (Float.min 2.9 (item [ 2 ] p)) in\n params := Nx.create f64 [| 3 |] [| item [ 0 ] p; z1; z2 |];\n state := s;\n let sigma_cur = item [] sigma_val in\n if sigma_cur < !best_sigma then begin\n best_sigma := sigma_cur;\n best_params := !params\n end;\n if i mod 50 = 0 || i = steps - 1 then begin\n let f_sky = 1.0 /. (1.0 +. Float.exp (-.item [ 0 ] !params)) in\n Printf.printf \"%5d %8.4f %8.3f %8.3f %10.6f\\n\" i f_sky\n (item [ 1 ] !params) (item [ 2 ] !params) sigma_cur\n end\n done;\n let f_sky_opt = 1.0 /. (1.0 +. Float.exp (-.item [ 0 ] !best_params)) in\n Printf.printf\n \"\\nGrad optimal: f_sky=%.4f bins=[0, %.2f, %.2f, 3.0] sigma(S8)=%.6f\\n\"\n f_sky_opt (item [ 1 ] !best_params) (item [ 2 ] !best_params) !best_sigma;\n\n (* Grid search validation *)\n let grid_best_sigma = ref infinity in\n let grid_best_fs = ref 0.0 in\n let grid_best_z1 = ref 0.0 in\n let grid_best_z2 = ref 0.0 in\n let n_fs = 12 and n_z1 = 15 and n_z2 = 15 in\n let n_grid_evals = ref 0 in\n Printf.printf \"\\nGrid search (%d*%d*%d)...\\n%!\" n_fs n_z1 n_z2;\n for fi = 0 to n_fs - 1 do\n let fs = 0.1 +. (Float.of_int fi *. 0.88 /. Float.of_int (n_fs - 1)) in\n let log_fs = Float.log (fs /. (1.0 -. fs)) in\n for z1i = 0 to n_z1 - 1 do\n let z1_v = 0.2 +. (Float.of_int z1i *. 2.4 /. Float.of_int (n_z1 - 1)) in\n for z2i = 0 to n_z2 - 1 do\n let z2_v =\n z1_v +. 0.15\n +. (Float.of_int z2i *. (2.7 -. z1_v) /. Float.of_int (n_z2 - 1))\n in\n if z2_v > z1_v +. 0.1 && z2_v < 2.9 then begin\n incr n_grid_evals;\n let p = Nx.create f64 [| 3 |] [| log_fs; z1_v; z2_v |] in\n let s = item [] (objective p) in\n if s < !grid_best_sigma then begin\n grid_best_sigma := s;\n grid_best_fs := fs;\n grid_best_z1 := z1_v;\n grid_best_z2 := z2_v\n end\n end\n done\n done\n done;\n Printf.printf\n \"Grid optimal: f_sky=%.4f bins=[0, %.2f, %.2f, 3.0] sigma(S8)=%.6f (%d \\\n evals)\\n\"\n !grid_best_fs !grid_best_z1 !grid_best_z2 !grid_best_sigma !n_grid_evals;\n Printf.printf \"\\nComparison:\\n\";\n Printf.printf \" Gradient: sigma(S8)=%.6f (%d evals)\\n\" !best_sigma steps;\n Printf.printf \" Grid: sigma(S8)=%.6f (%d evals)\\n\" !grid_best_sigma\n !n_grid_evals;\n let rel = (1.0 -. (!best_sigma /. !grid_best_sigma)) *. 100.0 in\n if rel >= 0.0 then\n Printf.printf \" Gradient %.1f%% better with %.0f* fewer evaluations\\n\" rel\n (Float.of_int !n_grid_evals /. Float.of_int steps)\n else\n Printf.printf\n \" Gradient within %.1f%% of grid with %.0f* fewer evaluations\\n\"\n (Float.abs rel)\n (Float.of_int !n_grid_evals /. Float.of_int steps)\n\nlet () =\n Printf.printf \"=== Differentiable Survey Optimization ===\\n\";\n Printf.printf \"Stage IV Weak Lensing Survey\\n\\n\";\n part1 ();\n part2 ()\n"
},
{
"path": "dev/umbra/examples/README.md",
"content": "# Umbra Examples\n\nLearn Umbra through progressively complex examples. Start with\n`01-constants-and-units` and work through the numbered examples in order.\n\n## Examples\n\n| Example | Concept | Key Functions |\n|---------|---------|---------------|\n| [`01-constants-and-units`](./01-constants-and-units/) | Type-safe physical quantities, conversions, constants | `Unit.Length.of_m`, `Const.c`, `Unit.Angle.deg` |\n| [`02-cosmological-distances`](./02-cosmological-distances/) | LCDM distances, SN Ia fitting | `Cosmo.luminosity_distance`, `Cosmo.distance_modulus` |\n| [`03-blackbody-fitting`](./03-blackbody-fitting/) | Fit stellar temperature from photometry | `Spectrum.blackbody`, `Photometry.ab_mag` |\n| [`04-extinction-and-magnitudes`](./04-extinction-and-magnitudes/) | Dust extinction, magnitude systems, K-corrections | `Extinction.ccm89`, `Photometry.vega_mag`, `Photometry.color` |\n| [`05-sed-fitting`](./05-sed-fitting/) | Full SED pipeline: blackbody, extinction, photometry | `Spectrum.blackbody`, `Extinction.apply`, `Photometry.ab_mag` |\n| [`06-coordinates-and-time`](./06-coordinates-and-time/) | Frame transforms, time scales, observer geometry | `Coord.galactic_of_icrs`, `Time.of_iso`, `Altaz.airmass` |\n| [`07-batch-photometry`](./07-batch-photometry/) | Batched operations over temperature and extinction grids | `Spectrum.blackbody`, `Extinction.apply`, `Photometry.ab_mag` |\n| [`08-photometric-redshifts`](./08-photometric-redshifts/) | Two-stage photo-z: grid search + gradient refinement | `Spectrum.redshift`, `Photometry.ab_mag`, `Rune.value_and_grad` |\n| [`09-gravitational-lensing`](./09-gravitational-lensing/) | Point-mass lens model parameter fitting | `Rune.value_and_grad`, `Vega.adam` |\n| [`10-uncertainty-propagation`](./10-uncertainty-propagation/) | AD Jacobians for error propagation vs Monte Carlo | `Rune.jacfwd`, `Cosmo.distance_modulus` |\n| [`11-bayesian-sed`](./11-bayesian-sed/) | Fisher matrix + HMC posterior sampling | `Rune.jacrev`, `Norn.hmc` |\n| [`12-survey-optimization`](./12-survey-optimization/) | Differentiable Fisher forecasting for survey design | `Survey.angular_cl`, `Cosmo.linear_power` |\n\n## Running Examples\n\nAll examples can be run with:\n\n```bash\ncd dev/umbra\ndune exec --root . examples//main.exe\n```\n\nFor example:\n\n```bash\ncd dev/umbra\ndune exec --root . examples/01-constants-and-units/main.exe\n```\n\n## Quick Reference\n\n### Cosmological Distances\n\n```ocaml\nopen Umbra\n\nlet cosmo = Cosmo.planck18 in\nlet z = Nx.scalar Nx.float64 0.5 in\nlet dl = Cosmo.luminosity_distance cosmo z\n```\n\n### Synthetic Photometry\n\n```ocaml\nlet sed =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 5800.0))\n ~wavelength:wave\n |> Extinction.apply (Extinction.ccm89 ~rv) ~av\n |> Spectrum.as_flux_density\nin\nlet mag = Photometry.ab_mag (Filters.sdss_r ()) sed\n```\n\n### Coordinate Transforms\n\n```ocaml\nlet ra = Unit.Angle.deg 83.633 in\nlet dec = Unit.Angle.deg (-5.550) in\nlet l, b = Coord.galactic_of_icrs ra dec\n```\n"
},
{
"path": "dev/umbra/lib/altaz.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet pi = Float.pi\n\ntype observer = { lat : float; lon : float; height : float }\n\nlet make_observer ~lat ~lon ?(height = Unit.Length.m 0.0) () =\n let lat = Nx.item [] (Unit.Angle.to_tensor lat) in\n let lon = Nx.item [] (Unit.Angle.to_tensor lon) in\n let height = Nx.item [] (Unit.Length.to_tensor height) in\n { lat; lon; height }\n\nlet observer_height obs =\n Unit.Length.of_tensor (Nx.scalar Nx.float64 obs.height)\n\ntype t = { az : Nx.float64_t; alt : Nx.float64_t }\n\nlet alt t = Unit.Angle.of_tensor t.alt\nlet az t = Unit.Angle.of_tensor t.az\n\n(* Earth Rotation Angle from UT1 Julian Date. ERA = 2π(0.7790572732640 +\n 1.00273781191135448 * Du) where Du = JD_UT1 - 2451545.0 *)\nlet era jd_ut1 =\n let du = jd_ut1 -. 2_451_545.0 in\n let theta =\n 2.0 *. pi *. (0.779_057_273_264_0 +. (1.002_737_811_911_354_48 *. du))\n in\n Float.rem theta (2.0 *. pi)\n\n(* IAU 2006 precession angles (Capitaine et al. 2003). T = Julian centuries from\n J2000.0 TT. Returns (zeta_A, z_A, theta_A) in radians. *)\nlet precession_angles t_cy =\n let arcsec_to_rad x = x *. pi /. 648_000.0 in\n let t2 = t_cy *. t_cy in\n let t3 = t2 *. t_cy in\n (* zeta_A = 2.5976176'' + 2306.0809506''T + 1.0109032''T² + 0.0182337''T³ *)\n let zeta_a =\n arcsec_to_rad\n (2.597_617_6 +. (2306.080_950_6 *. t_cy) +. (1.010_903_2 *. t2)\n +. (0.018_233_7 *. t3))\n in\n (* z_A = -2.5976176'' + 2306.0803226''T + 1.0947790''T² + 0.0182273''T³ *)\n let z_a =\n arcsec_to_rad\n (~-.2.597_617_6 +. (2306.080_322_6 *. t_cy) +. (1.094_779_0 *. t2)\n +. (0.018_227_3 *. t3))\n in\n (* theta_A = 2004.1917476''T - 0.4269353''T² - 0.0418251''T³ *)\n let theta_a =\n arcsec_to_rad\n ((2004.191_747_6 *. t_cy) -. (0.426_935_3 *. t2) -. (0.041_825_1 *. t3))\n in\n (zeta_a, z_a, theta_a)\n\n(* Apply IAU 2006 precession matrix to ICRS (RA, Dec) → mean (RA, Dec) of date.\n R = Rz(-z_A) · Ry(theta_A) · Rz(-zeta_A) *)\nlet precess_to_date ra dec t_cy =\n let zeta_a, z_a, theta_a = precession_angles t_cy in\n let sz = Float.sin zeta_a and cz = Float.cos zeta_a in\n let sa = Float.sin z_a and ca = Float.cos z_a in\n let st = Float.sin theta_a and ct = Float.cos theta_a in\n (* Rotation matrix elements *)\n let r11 = (ca *. ct *. cz) -. (sa *. sz) in\n let r12 = ~-.((ca *. ct *. sz) +. (sa *. cz)) in\n let r13 = ~-.(ca *. st) in\n let r21 = (sa *. ct *. cz) +. (ca *. sz) in\n let r22 = ~-.((sa *. ct *. sz) -. (ca *. cz)) in\n let r23 = ~-.(sa *. st) in\n let r31 = st *. cz in\n let r32 = ~-.(st *. sz) in\n let r33 = ct in\n let n = Nx.numel ra in\n let ra_out = Nx.zeros Nx.float64 [| n |] in\n let dec_out = Nx.zeros Nx.float64 [| n |] in\n for i = 0 to n - 1 do\n let r = Nx.item [ i ] ra in\n let d = Nx.item [ i ] dec in\n let cd = Float.cos d in\n let x = cd *. Float.cos r in\n let y = cd *. Float.sin r in\n let z = Float.sin d in\n let x' = (r11 *. x) +. (r12 *. y) +. (r13 *. z) in\n let y' = (r21 *. x) +. (r22 *. y) +. (r23 *. z) in\n let z' = (r31 *. x) +. (r32 *. y) +. (r33 *. z) in\n Nx.set_item [ i ] (Float.atan2 y' x') ra_out;\n Nx.set_item [ i ] (Float.asin (Float.max ~-.1.0 (Float.min 1.0 z'))) dec_out\n done;\n (ra_out, dec_out)\n\nlet airmass hz =\n let n = Nx.numel hz.alt in\n let out = Nx.zeros Nx.float64 [| n |] in\n let to_deg = 180.0 /. pi in\n for i = 0 to n - 1 do\n let alt_deg = Nx.item [ i ] hz.alt *. to_deg in\n (* Pickering (2002): X = 1 / sin(h + 244/(165 + 47h^1.1)) where h in deg *)\n let arg =\n alt_deg\n +. (244.0 /. (165.0 +. (47.0 *. Float.pow (Float.abs alt_deg) 1.1)))\n in\n let x = 1.0 /. Float.sin (arg *. pi /. 180.0) in\n Nx.set_item [ i ] (Float.max 1.0 x) out\n done;\n out\n\n(* Bennett (1982) atmospheric refraction for geometric altitude. R (arcmin) =\n cot(h + 7.31/(h + 4.4)) where h in degrees. Returns refraction in radians.\n Clamps to 0 below -1°. *)\nlet refraction_correction alt_rad =\n let h = alt_rad *. 180.0 /. pi in\n if h < -1.0 then 0.0\n else\n let arg = (h +. (7.31 /. (h +. 4.4))) *. pi /. 180.0 in\n let r_arcmin = 1.0 /. Float.tan arg in\n r_arcmin *. pi /. (180.0 *. 60.0)\n\nlet refraction hz =\n let n = Nx.numel hz.alt in\n let out = Nx.zeros Nx.float64 [| n |] in\n for i = 0 to n - 1 do\n Nx.set_item [ i ] (refraction_correction (Nx.item [ i ] hz.alt)) out\n done;\n Unit.Angle.of_tensor out\n\nlet of_coord ?(refraction = false) ~obstime ~observer c =\n let icrs = Coord.icrs c in\n let ra_rad = Unit.Angle.to_tensor (Coord.lon icrs) in\n let dec_rad = Unit.Angle.to_tensor (Coord.lat icrs) in\n (* Convert UTC → UT1 (ignoring DUT1 < 1s) then to TT for precession *)\n let jd_utc = Time.to_jd obstime in\n let jd_ut1 = jd_utc in\n let jd_tt = Time.to_jd (Time.tai_to_tt (Time.utc_to_tai obstime)) in\n let t_cy = (jd_tt -. 2_451_545.0) /. 36_525.0 in\n (* Precess ICRS to mean RA/Dec of date *)\n let ra_date, dec_date = precess_to_date ra_rad dec_rad t_cy in\n (* Hour angle: HA = ERA + observer_lon - RA_date *)\n let era_val = era jd_ut1 in\n let n = Nx.numel ra_rad in\n let alt_out = Nx.zeros Nx.float64 [| n |] in\n let az_out = Nx.zeros Nx.float64 [| n |] in\n let slat = Float.sin observer.lat and clat = Float.cos observer.lat in\n for i = 0 to n - 1 do\n let ha = era_val +. observer.lon -. Nx.item [ i ] ra_date in\n let dec = Nx.item [ i ] dec_date in\n let sdec = Float.sin dec and cdec = Float.cos dec in\n let sha = Float.sin ha and cha = Float.cos ha in\n (* alt = asin(sin(lat)sin(dec) + cos(lat)cos(dec)cos(ha)) *)\n let sin_alt = (slat *. sdec) +. (clat *. cdec *. cha) in\n let alt = Float.asin (Float.max ~-.1.0 (Float.min 1.0 sin_alt)) in\n (* az = atan2(-cos(dec)sin(ha), cos(lat)sin(dec) -\n sin(lat)cos(dec)cos(ha)) *)\n let num = ~-.(cdec *. sha) in\n let den = (clat *. sdec) -. (slat *. cdec *. cha) in\n let az = Float.atan2 num den in\n let az = if az < 0.0 then az +. (2.0 *. pi) else az in\n let alt = if refraction then alt +. refraction_correction alt else alt in\n Nx.set_item [ i ] alt alt_out;\n Nx.set_item [ i ] az az_out\n done;\n { alt = alt_out; az = az_out }\n\nlet to_coord ~obstime ~observer t =\n let jd_utc = Time.to_jd obstime in\n let jd_ut1 = jd_utc in\n let jd_tt = Time.to_jd (Time.tai_to_tt (Time.utc_to_tai obstime)) in\n let t_cy = (jd_tt -. 2_451_545.0) /. 36_525.0 in\n let era_val = era jd_ut1 in\n let slat = Float.sin observer.lat and clat = Float.cos observer.lat in\n let zeta_a, z_a, theta_a = precession_angles t_cy in\n (* Inverse precession matrix = transpose of forward *)\n let sz = Float.sin zeta_a and cz = Float.cos zeta_a in\n let sa = Float.sin z_a and ca = Float.cos z_a in\n let st = Float.sin theta_a and ct = Float.cos theta_a in\n let r11 = (ca *. ct *. cz) -. (sa *. sz) in\n let r12 = ~-.((ca *. ct *. sz) +. (sa *. cz)) in\n let r13 = ~-.(ca *. st) in\n let r21 = (sa *. ct *. cz) +. (ca *. sz) in\n let r22 = ~-.((sa *. ct *. sz) -. (ca *. cz)) in\n let r23 = ~-.(sa *. st) in\n let r31 = st *. cz in\n let r32 = ~-.(st *. sz) in\n let r33 = ct in\n (* Transpose for inverse *)\n let ri11 = r11 and ri12 = r21 and ri13 = r31 in\n let ri21 = r12 and ri22 = r22 and ri23 = r32 in\n let ri31 = r13 and ri32 = r23 and ri33 = r33 in\n let n = Nx.numel t.alt in\n let ra_out = Nx.zeros Nx.float64 [| n |] in\n let dec_out = Nx.zeros Nx.float64 [| n |] in\n for i = 0 to n - 1 do\n let alt = Nx.item [ i ] t.alt in\n let az = Nx.item [ i ] t.az in\n let salt = Float.sin alt and calt = Float.cos alt in\n let saz = Float.sin az and caz = Float.cos az in\n (* (Alt, Az) → (HA, Dec) *)\n let sin_dec = (slat *. salt) +. (clat *. calt *. caz) in\n let dec = Float.asin (Float.max ~-.1.0 (Float.min 1.0 sin_dec)) in\n let num = ~-.(calt *. saz) in\n let den = (clat *. salt) -. (slat *. calt *. caz) in\n let ha = Float.atan2 num den in\n (* RA_date = ERA + observer_lon - HA *)\n let ra_date = era_val +. observer.lon -. ha in\n (* Deprecess: mean of date → ICRS *)\n let cd = Float.cos dec in\n let x = cd *. Float.cos ra_date in\n let y = cd *. Float.sin ra_date in\n let z = Float.sin dec in\n let x' = (ri11 *. x) +. (ri12 *. y) +. (ri13 *. z) in\n let y' = (ri21 *. x) +. (ri22 *. y) +. (ri23 *. z) in\n let z' = (ri31 *. x) +. (ri32 *. y) +. (ri33 *. z) in\n let ra = Float.atan2 y' x' in\n let ra = if ra < 0.0 then ra +. (2.0 *. pi) else ra in\n let dec = Float.asin (Float.max ~-.1.0 (Float.min 1.0 z')) in\n Nx.set_item [ i ] ra ra_out;\n Nx.set_item [ i ] dec dec_out\n done;\n Coord.of_radec\n ~ra:(Unit.Angle.of_tensor ra_out)\n ~dec:(Unit.Angle.of_tensor dec_out)\n"
},
{
"path": "dev/umbra/lib/altaz.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Altitude-azimuth (horizontal) coordinates.\n\n Converts celestial coordinates to local horizon coordinates for a given\n observer location and time. Uses IAU 2006 precession (Capitaine et al. 2003)\n and the Earth Rotation Angle.\n\n {b Warning.} Nutation and polar motion are omitted. Atmospheric refraction\n can be applied via {!refraction} or the [~refraction] parameter of\n {!of_coord}. Accuracy is ~1 arcminute for dates within a few centuries of\n J2000.0.\n\n {[\n let obs = Altaz.make_observer ~lat:(Unit.Angle.deg 28.7624) ~lon:(Unit.Angle.deg (-17.8792)) () in\n let t = Time.of_iso \"2024-06-21T22:00:00\" in\n let vega =\n Coord.of_radec\n ~ra:(Unit.Angle.deg 279.2347)\n ~dec:(Unit.Angle.deg 38.7837)\n in\n let hz = Altaz.of_coord ~obstime:t ~observer:obs vega in\n let alt_deg = Nx.item [] (Unit.Angle.in_deg (Altaz.alt hz))\n ]} *)\n\n(** {1:observer Observer} *)\n\ntype observer\n(** The type for a ground-based observer location. *)\n\nval make_observer :\n lat:Unit.angle Unit.t ->\n lon:Unit.angle Unit.t ->\n ?height:Unit.length Unit.t ->\n unit ->\n observer\n(** [make_observer ~lat ~lon ?height ()] is an observer at geodetic latitude\n [lat], longitude [lon], and elevation [height] above the reference\n ellipsoid. [lon] is positive East. [height] defaults to sea level.\n\n [height] is stored for forward compatibility but does not yet affect\n coordinate transforms. *)\n\nval observer_height : observer -> Unit.length Unit.t\n(** [observer_height obs] is the observer's elevation above the reference\n ellipsoid. *)\n\n(** {1:coords Horizontal coordinates} *)\n\ntype t\n(** The type for altitude-azimuth coordinates. Azimuth is measured from North\n through East. *)\n\nval alt : t -> Unit.angle Unit.t\n(** [alt t] is the altitude (elevation above the horizon). *)\n\nval az : t -> Unit.angle Unit.t\n(** [az t] is the azimuth (North = 0, East = 90 deg). *)\n\n(** {1:derived Derived quantities} *)\n\nval airmass : t -> Nx.float64_t\n(** [airmass hz] is the airmass at the altitude of [hz] using the Pickering\n (2002) formula. Well-behaved from zenith to horizon. Not differentiable\n (operates on float-level altitude values). *)\n\n(** {1:refraction Atmospheric refraction} *)\n\nval refraction : t -> Unit.angle Unit.t\n(** [refraction hz] is the atmospheric refraction correction at the geometric\n altitude of [hz], using the Bennett (1982) formula. The correction is\n positive (objects appear higher than their geometric position). Returns zero\n for altitudes below -1°.\n\n Not differentiable (scalar-level trigonometry). *)\n\n(** {1:converting Converting} *)\n\nval of_coord :\n ?refraction:bool ->\n obstime:Time.utc Time.t ->\n observer:observer ->\n Coord.t ->\n t\n(** [of_coord ~obstime ~observer c] converts celestial coordinates [c] to\n altitude-azimuth for [observer] at [obstime]. Applies IAU 2006 precession to\n move from ICRS to the mean equator of date.\n\n When [refraction] is [true], the Bennett (1982) atmospheric refraction\n correction is applied to the computed altitude. [refraction] defaults to\n [false].\n\n Not differentiable (scalar-level rotation matrices). *)\n\nval to_coord : obstime:Time.utc Time.t -> observer:observer -> t -> Coord.t\n(** [to_coord ~obstime ~observer t] converts altitude-azimuth coordinates [t]\n back to ICRS celestial coordinates for [observer] at [obstime]. Not\n differentiable (scalar-level rotation matrices). *)\n"
},
{
"path": "dev/umbra/lib/const.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Fundamental constants (CODATA 2022) *)\n\nlet c = Unit.Velocity.m_s 299_792_458.0\nlet m_e = Unit.Mass.kg 9.109_383_713_9e-31\nlet m_p = Unit.Mass.kg 1.672_621_923_69e-27\nlet m_n = Unit.Mass.kg 1.674_927_498_04e-27\nlet u = Unit.Mass.kg 1.660_539_066_60e-27\n\n(* Astronomical constants (IAU 2015) *)\n\nlet au = Unit.Length.au 1.0\nlet pc = Unit.Length.pc 1.0\nlet solar_mass = Unit.Mass.solar_mass 1.0\nlet solar_radius = Unit.Length.solar_radius 1.0\nlet solar_luminosity = Unit.Power.solar_luminosity 1.0\nlet earth_mass = Unit.Mass.earth_mass 1.0\nlet earth_radius = Unit.Length.earth_radius 1.0\nlet jupiter_mass = Unit.Mass.jupiter_mass 1.0\nlet jupiter_radius = Unit.Length.jupiter_radius 1.0\n\n(* Raw SI floats for compound dimensions (CODATA 2022) *)\n\nlet h_si = 6.626_070_15e-34\nlet hbar_si = 1.054_571_817e-34\nlet g_si = 6.674_30e-11\nlet k_b_si = 1.380_649e-23\nlet sigma_sb_si = 5.670_374_419e-8\nlet n_a = 6.022_140_76e23\nlet sigma_t_si = 6.652_458_705_1e-29\nlet b_wien_si = 2.897_771_955e-3\nlet alpha = 7.297_352_5643e-3\nlet a_0 = Unit.Length.m 5.291_772_105_44e-11\nlet gm_sun_si = 1.327_124_4e20\nlet gm_earth_si = 3.986_004e14\nlet gm_jup_si = 1.266_865_3e17\nlet l_bol0 = Unit.Power.w 3.0128e28\n"
},
{
"path": "dev/umbra/lib/const.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Physical and astronomical constants.\n\n Typed constants use {!Unit.t} with the appropriate phantom dimension. Raw SI\n floats are provided for compound dimensions that do not map to a single\n {!Unit} dimension type.\n\n Fundamental constants follow\n {{:https://physics.nist.gov/cuu/Constants/}CODATA 2022}. Astronomical\n constants follow IAU 2015. *)\n\n(** {1:fundamental Fundamental constants} *)\n\nval c : Unit.velocity Unit.t\n(** [c] is the speed of light in vacuum (299 792 458 m/s, exact). *)\n\n(** {1:particle Particle masses} *)\n\nval m_e : Unit.mass Unit.t\n(** [m_e] is the electron mass (9.109 383 7139e-31 kg). *)\n\nval m_p : Unit.mass Unit.t\n(** [m_p] is the proton mass (1.672 621 923 69e-27 kg). *)\n\nval m_n : Unit.mass Unit.t\n(** [m_n] is the neutron mass (1.674 927 498 04e-27 kg). *)\n\nval u : Unit.mass Unit.t\n(** [u] is the atomic mass unit (1.660 539 066 60e-27 kg). *)\n\n(** {1:astro Astronomical constants} *)\n\nval au : Unit.length Unit.t\n(** [au] is one astronomical unit. *)\n\nval pc : Unit.length Unit.t\n(** [pc] is one parsec. *)\n\nval solar_mass : Unit.mass Unit.t\n(** [solar_mass] is one solar mass. *)\n\nval solar_radius : Unit.length Unit.t\n(** [solar_radius] is one solar radius. *)\n\nval solar_luminosity : Unit.power Unit.t\n(** [solar_luminosity] is one solar luminosity. *)\n\nval earth_mass : Unit.mass Unit.t\n(** [earth_mass] is one Earth mass. *)\n\nval earth_radius : Unit.length Unit.t\n(** [earth_radius] is one Earth radius. *)\n\nval jupiter_mass : Unit.mass Unit.t\n(** [jupiter_mass] is one Jupiter mass. *)\n\nval jupiter_radius : Unit.length Unit.t\n(** [jupiter_radius] is one Jupiter radius. *)\n\n(** {1:si Raw SI constants}\n\n Constants with compound dimensions that do not map to a single {!Unit}\n dimension type. CODATA 2022 values. *)\n\nval h_si : float\n(** [h_si] is the Planck constant (6.626 070 15e-34 J s, exact). *)\n\nval hbar_si : float\n(** [hbar_si] is the reduced Planck constant (1.054 571 817e-34 J s). *)\n\nval g_si : float\n(** [g_si] is the gravitational constant (6.674 30e-11 m{^ 3} kg{^ -1} s{^ -2}).\n*)\n\nval k_b_si : float\n(** [k_b_si] is the Boltzmann constant (1.380 649e-23 J K{^ -1}, exact). *)\n\nval sigma_sb_si : float\n(** [sigma_sb_si] is the Stefan-Boltzmann constant (5.670 374 419e-8 W m{^ -2}\n K{^ -4}). *)\n\nval n_a : float\n(** [n_a] is the Avogadro constant (6.022 140 76e23 mol{^ -1}, exact). *)\n\nval sigma_t_si : float\n(** [sigma_t_si] is the Thomson scattering cross-section (6.652 458 705 1e-29\n m{^ 2}). *)\n\nval b_wien_si : float\n(** [b_wien_si] is the Wien displacement law constant (2.897 771 955e-3 m K). *)\n\nval alpha : float\n(** [alpha] is the fine-structure constant (7.297 352 5643e-3). *)\n\nval a_0 : Unit.length Unit.t\n(** [a_0] is the Bohr radius (5.291 772 105 44e-11 m). *)\n\nval gm_sun_si : float\n(** [gm_sun_si] is the solar mass parameter (1.327 124 4e20 m{^ 3} s{^ -2}).\n More precise than [g_si * solar_mass] for orbital mechanics. *)\n\nval gm_earth_si : float\n(** [gm_earth_si] is the Earth mass parameter (3.986 004e14 m{^ 3} s{^ -2}). *)\n\nval gm_jup_si : float\n(** [gm_jup_si] is the Jupiter mass parameter (1.266 865 3e17 m{^ 3} s{^ -2}).\n*)\n\nval l_bol0 : Unit.power Unit.t\n(** [l_bol0] is the IAU 2015 zero-point bolometric luminosity (3.0128e28 W). *)\n"
},
{
"path": "dev/umbra/lib/coord.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet pi = Float.pi\nlet deg_to_rad = pi /. 180.0\nlet two_pi = Nx.scalar Nx.float64 (2.0 *. pi)\n\ntype frame = ICRS | Galactic | Ecliptic_j2000 | Supergalactic\n\n(* Internally stores lon/lat in radians *)\ntype t = { frame : frame; lon : Nx.float64_t; lat : Nx.float64_t }\n\n(* IAU rotation matrices *)\n\nlet ra_gp = 192.85948 *. deg_to_rad\nlet dec_gp = 27.12825 *. deg_to_rad\nlet l_ncp = 122.93192 *. deg_to_rad\n\nlet icrs_to_gal =\n let sd = Float.sin dec_gp and cd = Float.cos dec_gp in\n let sa = Float.sin ra_gp and ca = Float.cos ra_gp in\n let sl = Float.sin l_ncp and cl = Float.cos l_ncp in\n [|\n [|\n (~-.sl *. sa) -. (cl *. ca *. sd);\n (sl *. ca) -. (cl *. sa *. sd);\n cl *. cd;\n |];\n [|\n (cl *. sa) -. (sl *. ca *. sd);\n (~-.cl *. ca) -. (sl *. sa *. sd);\n sl *. cd;\n |];\n [| ca *. cd; sa *. cd; sd |];\n |]\n\nlet transpose_3x3 m =\n [|\n [| m.(0).(0); m.(1).(0); m.(2).(0) |];\n [| m.(0).(1); m.(1).(1); m.(2).(1) |];\n [| m.(0).(2); m.(1).(2); m.(2).(2) |];\n |]\n\nlet gal_to_icrs = transpose_3x3 icrs_to_gal\n\n(* Fixed J2000.0 mean obliquity: 23.4392911 degrees *)\nlet obliquity = 23.4392911 *. deg_to_rad\n\nlet icrs_to_ecl =\n let se = Float.sin obliquity and ce = Float.cos obliquity in\n [| [| 1.0; 0.0; 0.0 |]; [| 0.0; ce; se |]; [| 0.0; ~-.se; ce |] |]\n\nlet ecl_to_icrs = transpose_3x3 icrs_to_ecl\n\n(* Supergalactic: defined relative to Galactic. SGP at (l=47.37, b=6.32), SGL\n origin at l=137.37 *)\nlet sgl_l0 = 137.37 *. deg_to_rad\nlet sgp_l = 47.37 *. deg_to_rad\nlet sgp_b = 6.32 *. deg_to_rad\n\nlet gal_to_sgal =\n let sb = Float.sin sgp_b and cb = Float.cos sgp_b in\n let sl = Float.sin sgp_l and cl = Float.cos sgp_l in\n let sl0 = Float.sin sgl_l0 and cl0 = Float.cos sgl_l0 in\n let r00 = (~-.sl0 *. sl) -. (cl0 *. cl *. sb) in\n let r01 = (sl0 *. cl) -. (cl0 *. sl *. sb) in\n let r02 = cl0 *. cb in\n let r10 = (cl0 *. sl) -. (sl0 *. cl *. sb) in\n let r11 = (~-.cl0 *. cl) -. (sl0 *. sl *. sb) in\n let r12 = sl0 *. cb in\n let r20 = cl *. cb in\n let r21 = sl *. cb in\n let r22 = sb in\n [| [| r00; r01; r02 |]; [| r10; r11; r12 |]; [| r20; r21; r22 |] |]\n\nlet sgal_to_gal = transpose_3x3 gal_to_sgal\n\nlet rotate mat lon_rad lat_rad =\n let cl = Nx.cos lat_rad and sl = Nx.sin lat_rad in\n let ca = Nx.cos lon_rad and sa = Nx.sin lon_rad in\n let x = Nx.mul cl ca and y = Nx.mul cl sa in\n let x' =\n Nx.add\n (Nx.add (Nx.mul_s x mat.(0).(0)) (Nx.mul_s y mat.(0).(1)))\n (Nx.mul_s sl mat.(0).(2))\n in\n let y' =\n Nx.add\n (Nx.add (Nx.mul_s x mat.(1).(0)) (Nx.mul_s y mat.(1).(1)))\n (Nx.mul_s sl mat.(1).(2))\n in\n let z' =\n Nx.add\n (Nx.add (Nx.mul_s x mat.(2).(0)) (Nx.mul_s y mat.(2).(1)))\n (Nx.mul_s sl mat.(2).(2))\n in\n let z_clamped = Nx.clamp ~min:(-1.0) ~max:1.0 z' in\n let lat' = Nx.asin z_clamped in\n let lon' = Nx.atan2 y' x' in\n let mask = Nx.less_s lon' 0.0 in\n let lon' = Nx.where mask (Nx.add lon' two_pi) lon' in\n (lon', lat')\n\nlet ensure_1d t = if Nx.ndim t = 0 then Nx.reshape [| 1 |] t else t\n\nlet make frame ~lon ~lat =\n let lon_rad = ensure_1d (Unit.Angle.to_tensor lon) in\n let lat_rad = ensure_1d (Unit.Angle.to_tensor lat) in\n if Nx.ndim lon_rad <> 1 || Nx.ndim lat_rad <> 1 then\n invalid_arg \"Coord: lon and lat must be scalar or 1-D tensors\";\n if Nx.numel lon_rad <> Nx.numel lat_rad then\n invalid_arg \"Coord: lon and lat must have the same length\";\n { frame; lon = lon_rad; lat = lat_rad }\n\nlet of_radec ~ra ~dec = make ICRS ~lon:ra ~lat:dec\nlet of_galactic ~l ~b = make Galactic ~lon:l ~lat:b\nlet of_ecliptic_j2000 ~lon ~lat = make Ecliptic_j2000 ~lon ~lat\nlet of_supergalactic ~sgl ~sgb = make Supergalactic ~lon:sgl ~lat:sgb\nlet frame c = c.frame\nlet size c = Nx.numel c.lon\nlet lon c = Unit.Angle.of_tensor c.lon\nlet lat c = Unit.Angle.of_tensor c.lat\n\nlet to_icrs c =\n match c.frame with\n | ICRS -> c\n | Galactic ->\n let lon', lat' = rotate gal_to_icrs c.lon c.lat in\n { frame = ICRS; lon = lon'; lat = lat' }\n | Ecliptic_j2000 ->\n let lon', lat' = rotate ecl_to_icrs c.lon c.lat in\n { frame = ICRS; lon = lon'; lat = lat' }\n | Supergalactic ->\n let gal_lon, gal_lat = rotate sgal_to_gal c.lon c.lat in\n let icrs_lon, icrs_lat = rotate gal_to_icrs gal_lon gal_lat in\n { frame = ICRS; lon = icrs_lon; lat = icrs_lat }\n\nlet ra c = lon (to_icrs c)\nlet dec c = lat (to_icrs c)\n\nlet to_frame target c =\n if c.frame = target then c\n else\n let icrs = to_icrs c in\n match target with\n | ICRS -> icrs\n | Galactic ->\n let lon', lat' = rotate icrs_to_gal icrs.lon icrs.lat in\n { frame = Galactic; lon = lon'; lat = lat' }\n | Ecliptic_j2000 ->\n let lon', lat' = rotate icrs_to_ecl icrs.lon icrs.lat in\n { frame = Ecliptic_j2000; lon = lon'; lat = lat' }\n | Supergalactic ->\n let gal_lon, gal_lat = rotate icrs_to_gal icrs.lon icrs.lat in\n let sg_lon, sg_lat = rotate gal_to_sgal gal_lon gal_lat in\n { frame = Supergalactic; lon = sg_lon; lat = sg_lat }\n\nlet icrs c = to_frame ICRS c\nlet galactic c = to_frame Galactic c\nlet ecliptic_j2000 c = to_frame Ecliptic_j2000 c\nlet supergalactic c = to_frame Supergalactic c\n\nlet trig_of a b =\n let a = to_icrs a and b = to_icrs b in\n let dlon = Nx.sub b.lon a.lon in\n let cos_lat1 = Nx.cos a.lat and sin_lat1 = Nx.sin a.lat in\n let cos_lat2 = Nx.cos b.lat and sin_lat2 = Nx.sin b.lat in\n let cos_dlon = Nx.cos dlon and sin_dlon = Nx.sin dlon in\n (dlon, cos_lat1, sin_lat1, cos_lat2, sin_lat2, cos_dlon, sin_dlon)\n\nlet separation a b =\n if size a <> size b then\n invalid_arg \"Coord.separation: arrays must have the same length\";\n let _, cos_lat1, sin_lat1, cos_lat2, sin_lat2, cos_dlon, sin_dlon =\n trig_of a b\n in\n (* Vincenty formula *)\n let a1 = Nx.mul cos_lat2 sin_dlon in\n let a2 =\n Nx.sub (Nx.mul cos_lat1 sin_lat2)\n (Nx.mul (Nx.mul sin_lat1 cos_lat2) cos_dlon)\n in\n let num = Nx.sqrt (Nx.add (Nx.square a1) (Nx.square a2)) in\n let den =\n Nx.add (Nx.mul sin_lat1 sin_lat2)\n (Nx.mul (Nx.mul cos_lat1 cos_lat2) cos_dlon)\n in\n let sep = Nx.atan2 num den in\n Unit.Angle.of_tensor (Nx.abs sep)\n\nlet position_angle a b =\n if size a <> size b then\n invalid_arg \"Coord.position_angle: arrays must have the same length\";\n let _, cos_lat1, sin_lat1, cos_lat2, sin_lat2, cos_dlon, sin_dlon =\n trig_of a b\n in\n let num = Nx.mul cos_lat2 sin_dlon in\n let den =\n Nx.sub (Nx.mul cos_lat1 sin_lat2)\n (Nx.mul (Nx.mul sin_lat1 cos_lat2) cos_dlon)\n in\n let pa = Nx.atan2 num den in\n let mask = Nx.less_s pa 0.0 in\n Unit.Angle.of_tensor (Nx.where mask (Nx.add pa two_pi) pa)\n\n(* --- Offset operations --- *)\n\nlet offset_by ~position_angle ~separation c =\n let pa = Unit.Angle.to_tensor position_angle in\n let sep = Unit.Angle.to_tensor separation in\n let cos_sep = Nx.cos sep and sin_sep = Nx.sin sep in\n let cos_pa = Nx.cos pa and sin_pa = Nx.sin pa in\n let sin_lat = Nx.sin c.lat and cos_lat = Nx.cos c.lat in\n (* lat2 = asin(sin(lat1)*cos(sep) + cos(lat1)*sin(sep)*cos(pa)) *)\n let sin_lat2 =\n Nx.add (Nx.mul sin_lat cos_sep) (Nx.mul (Nx.mul cos_lat sin_sep) cos_pa)\n in\n let lat2 = Nx.asin (Nx.clamp ~min:(-1.0) ~max:1.0 sin_lat2) in\n (* lon2 = lon1 + atan2(sin(pa)*sin(sep), cos(lat1)*cos(sep) -\n sin(lat1)*sin(sep)*cos(pa)) *)\n let num = Nx.mul sin_pa sin_sep in\n let den =\n Nx.sub (Nx.mul cos_lat cos_sep) (Nx.mul (Nx.mul sin_lat sin_sep) cos_pa)\n in\n let dlon = Nx.atan2 num den in\n let lon2 = Nx.add c.lon dlon in\n let lon2 = Nx.where (Nx.less_s lon2 0.0) (Nx.add lon2 two_pi) lon2 in\n let lon2 =\n Nx.where (Nx.greater_equal lon2 two_pi) (Nx.sub lon2 two_pi) lon2\n in\n { frame = c.frame; lon = lon2; lat = lat2 }\n\nlet spherical_offsets_to a b =\n if size a <> size b then\n invalid_arg \"Coord.spherical_offsets_to: arrays must have the same length\";\n if a.frame <> b.frame then\n invalid_arg\n \"Coord.spherical_offsets_to: coordinates must be in the same frame\";\n (* Δlon = (lon_b - lon_a) * cos(lat_a), Δlat = lat_b - lat_a *)\n let dlon = Nx.mul (Nx.sub b.lon a.lon) (Nx.cos a.lat) in\n let dlat = Nx.sub b.lat a.lat in\n (Unit.Angle.of_tensor dlon, Unit.Angle.of_tensor dlat)\n\n(* --- Catalog cross-matching --- *)\n\ntype coord = t\ntype result = { indices : Nx.int32_t; separations : Unit.angle Unit.t }\n\ntype within_result = {\n indices_a : Nx.int32_t;\n indices_b : Nx.int32_t;\n separations : Unit.angle Unit.t;\n}\n\nlet to_xyz c =\n let icrs = to_icrs c in\n let n = size c in\n let xs = Array.make n 0.0 in\n let ys = Array.make n 0.0 in\n let zs = Array.make n 0.0 in\n for i = 0 to n - 1 do\n let r = Nx.item [ i ] icrs.lon in\n let d = Nx.item [ i ] icrs.lat in\n let cd = Float.cos d in\n xs.(i) <- cd *. Float.cos r;\n ys.(i) <- cd *. Float.sin r;\n zs.(i) <- Float.sin d\n done;\n (xs, ys, zs)\n\nlet chord_to_rad chord_sq =\n let chord = Float.sqrt (Float.max 0.0 chord_sq) in\n let half_chord = Float.min 1.0 (chord /. 2.0) in\n 2.0 *. Float.asin half_chord\n\nmodule Index = struct\n type t = { tree : Kdtree.t }\n\n let of_coord c =\n let xs, ys, zs = to_xyz c in\n let tree = Kdtree.build xs ys zs in\n { tree }\n\n let nearest idx query =\n let qx, qy, qz = to_xyz query in\n let n = Array.length qx in\n let indices = Nx.zeros Nx.int32 [| n |] in\n let seps = Nx.zeros Nx.float64 [| n |] in\n for i = 0 to n - 1 do\n let j, dist_sq = Kdtree.nearest idx.tree qx.(i) qy.(i) qz.(i) in\n Nx.set_item [ i ] (Int32.of_int j) indices;\n Nx.set_item [ i ] (chord_to_rad dist_sq) seps\n done;\n { indices; separations = Unit.Angle.of_tensor seps }\n\n let within idx query ~max_sep =\n let max_sep_rad = Nx.item [] (Unit.Angle.to_tensor max_sep) in\n let half_angle = max_sep_rad /. 2.0 in\n let chord = 2.0 *. Float.sin half_angle in\n let max_dist_sq = chord *. chord in\n let qx, qy, qz = to_xyz query in\n let na = Array.length qx in\n let acc = ref [] and count = ref 0 in\n for i = 0 to na - 1 do\n let matches = Kdtree.within idx.tree qx.(i) qy.(i) qz.(i) max_dist_sq in\n List.iter\n (fun (j, dist_sq) ->\n acc := (i, j, chord_to_rad dist_sq) :: !acc;\n incr count)\n matches\n done;\n let n = !count in\n let out_a = Nx.zeros Nx.int32 [| n |] in\n let out_b = Nx.zeros Nx.int32 [| n |] in\n let out_s = Nx.zeros Nx.float64 [| n |] in\n let k = ref (n - 1) in\n List.iter\n (fun (i, j, sep) ->\n let k' = !k in\n Nx.set_item [ k' ] (Int32.of_int i) out_a;\n Nx.set_item [ k' ] (Int32.of_int j) out_b;\n Nx.set_item [ k' ] sep out_s;\n decr k)\n !acc;\n {\n indices_a = out_a;\n indices_b = out_b;\n separations = Unit.Angle.of_tensor out_s;\n }\nend\n\nlet nearest query catalog =\n if size catalog = 0 then invalid_arg \"Coord.nearest: catalog is empty\";\n let idx = Index.of_coord catalog in\n Index.nearest idx query\n\nlet within a b ~max_sep =\n let idx = Index.of_coord b in\n Index.within idx a ~max_sep\n"
},
{
"path": "dev/umbra/lib/coord.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Celestial coordinates with frame transforms and catalog matching.\n\n Positions are stored as longitude/latitude pairs in 1D {!Unit.angle}\n quantities and can be converted between {!ICRS}, {!Galactic},\n {!Ecliptic_j2000}, and {!Supergalactic} frames via 3x3 rotation matrices.\n\n {[\n let c = Coord.of_radec ~ra:(Unit.Angle.of_deg ra) ~dec:(Unit.Angle.of_deg dec) in\n let gal = Coord.galactic c\n ]} *)\n\n(** {1:types Types} *)\n\n(** The type for celestial reference frames. *)\ntype frame =\n | ICRS (** International Celestial Reference System. *)\n | Galactic (** IAU Galactic coordinates. *)\n | Ecliptic_j2000 (** Ecliptic coordinates at J2000.0 epoch. *)\n | Supergalactic (** Supergalactic coordinates. *)\n\ntype t\n(** The type for celestial coordinates. A pair of 1D angle quantities\n (longitude, latitude), tagged with a {!frame}. *)\n\n(** {1:constructors Constructors}\n\n All constructors require 1D angle quantities of equal length.\n\n Raises [Invalid_argument] if the tensors are not 1D or differ in length. *)\n\nval of_radec : ra:Unit.angle Unit.t -> dec:Unit.angle Unit.t -> t\n(** [of_radec ~ra ~dec] is a coordinate in the ICRS frame. [ra] and [dec] must\n be scalar or 1-D angle quantities with matching sizes. *)\n\nval of_galactic : l:Unit.angle Unit.t -> b:Unit.angle Unit.t -> t\n(** [of_galactic ~l ~b] is a coordinate in the Galactic frame. [l] and [b] must\n be scalar or 1-D angle quantities with matching sizes. *)\n\nval of_ecliptic_j2000 : lon:Unit.angle Unit.t -> lat:Unit.angle Unit.t -> t\n(** [of_ecliptic_j2000 ~lon ~lat] is a coordinate in the ecliptic frame at the\n J2000.0 mean obliquity (23.4392911 degrees). [lon] and [lat] must be scalar\n or 1-D angle quantities with matching sizes. *)\n\nval of_supergalactic : sgl:Unit.angle Unit.t -> sgb:Unit.angle Unit.t -> t\n(** [of_supergalactic ~sgl ~sgb] is a coordinate in the Supergalactic frame.\n [sgl] and [sgb] must be scalar or 1-D angle quantities with matching sizes.\n*)\n\n(** {1:accessors Accessors} *)\n\nval frame : t -> frame\n(** [frame c] is the reference frame of [c]. *)\n\nval size : t -> int\n(** [size c] is the number of positions in [c]. *)\n\nval lon : t -> Unit.angle Unit.t\n(** [lon c] is the longitude component of [c]. *)\n\nval lat : t -> Unit.angle Unit.t\n(** [lat c] is the latitude component of [c]. *)\n\nval ra : t -> Unit.angle Unit.t\n(** [ra c] is the ICRS right ascension of [c]. Converts to ICRS first if [c] is\n in another frame. *)\n\nval dec : t -> Unit.angle Unit.t\n(** [dec c] is the ICRS declination of [c]. Converts to ICRS first if [c] is in\n another frame. *)\n\n(** {1:transforms Frame transforms} *)\n\nval to_frame : frame -> t -> t\n(** [to_frame f c] is [c] converted to frame [f]. Returns [c] unchanged if [c]\n is already in [f]. All conversions go through ICRS as the pivot frame. Not\n differentiable (scalar-level rotation matrices). *)\n\nval icrs : t -> t\n(** [icrs c] is [to_frame ICRS c]. *)\n\nval galactic : t -> t\n(** [galactic c] is [to_frame Galactic c]. *)\n\nval ecliptic_j2000 : t -> t\n(** [ecliptic_j2000 c] is [to_frame Ecliptic_j2000 c]. *)\n\nval supergalactic : t -> t\n(** [supergalactic c] is [to_frame Supergalactic c]. *)\n\n(** {1:separation Angular separation} *)\n\nval separation : t -> t -> Unit.angle Unit.t\n(** [separation a b] is the angular separation between corresponding positions\n of [a] and [b], computed with the Vincenty formula. Both coordinates are\n converted to ICRS before computation. Not differentiable (scalar-level\n trigonometry).\n\n Raises [Invalid_argument] if [a] and [b] differ in {!size}. *)\n\nval position_angle : t -> t -> Unit.angle Unit.t\n(** [position_angle a b] is the position angle from [a] to [b], measured North\n through East, in \\[0, 2{e pi}). Both coordinates are converted to ICRS\n before computation. Not differentiable (scalar-level trigonometry).\n\n Raises [Invalid_argument] if [a] and [b] differ in {!size}. *)\n\n(** {1:offsets Offset operations} *)\n\nval offset_by :\n position_angle:Unit.angle Unit.t -> separation:Unit.angle Unit.t -> t -> t\n(** [offset_by ~position_angle ~separation c] is the coordinate obtained by\n moving each position in [c] along bearing [position_angle] (North through\n East) by angular distance [separation]. The result is in the same frame as\n [c]. Not differentiable (scalar-level trigonometry). *)\n\nval spherical_offsets_to : t -> t -> Unit.angle Unit.t * Unit.angle Unit.t\n(** [spherical_offsets_to a b] is [(dlon, dlat)] where\n [dlon = (lon_b - lon_a) * cos(lat_a)] and [dlat = lat_b - lat_a]. Both\n coordinates must be in the same frame. Not differentiable (scalar-level\n trigonometry).\n\n Raises [Invalid_argument] if [a] and [b] differ in {!size} or {!frame}. *)\n\n(** {1:matching Catalog cross-matching}\n\n Matches positions between catalogs using a 3D kd-tree built from unit-sphere\n Cartesian coordinates. All indices in results are 0-based.\n\n {b Warning.} Cross-matching is not differentiable: it produces integer\n indices and uses discrete tree search. *)\n\ntype coord = t\n(** Alias for {!t}, used inside {!Index} to avoid shadowing. *)\n\ntype result = {\n indices : Nx.int32_t; (** 0-based indices into the catalog. *)\n separations : Unit.angle Unit.t; (** Angular distances. *)\n}\n(** The type for nearest-match results. For each query position, {!indices}\n gives the index of the nearest catalog entry and {!separations} gives the\n angular distance to it. Both have the same length as the query. *)\n\ntype within_result = {\n indices_a : Nx.int32_t; (** 0-based indices into the query. *)\n indices_b : Nx.int32_t; (** 0-based indices into the catalog. *)\n separations : Unit.angle Unit.t; (** Angular distances. *)\n}\n(** The type for within-radius match results. Each entry represents one matched\n pair. The three fields have equal length. *)\n\n(** {2:index Reusable index}\n\n Build a kd-tree once and query it many times. *)\n\nmodule Index : sig\n type t\n (** The type for a prebuilt spatial index over a catalog. *)\n\n val of_coord : coord -> t\n (** [of_coord c] builds a kd-tree index from the positions in [c]. Coordinates\n are converted to ICRS internally. *)\n\n val nearest : t -> coord -> result\n (** [nearest idx query] finds, for each position in [query], the nearest\n position in the indexed catalog. *)\n\n val within : t -> coord -> max_sep:Unit.angle Unit.t -> within_result\n (** [within idx query ~max_sep] finds all pairs where a position in [query] is\n within [max_sep] of a position in the indexed catalog. *)\nend\n\nval nearest : t -> t -> result\n(** [nearest query catalog] finds, for each position in [query], the nearest\n position in [catalog].\n\n Raises [Invalid_argument] if [catalog] is empty. *)\n\nval within : t -> t -> max_sep:Unit.angle Unit.t -> within_result\n(** [within a b ~max_sep] finds all pairs of positions where the separation is\n at most [max_sep]. Builds a kd-tree on [b]. *)\n"
},
{
"path": "dev/umbra/lib/cosmo.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Cosmological distance calculations for ΛCDM, wCDM, and w0waCDM universes.\n\n w0waCDM subsumes all models: - flat ΛCDM: omega_k = 0, w0 = -1, wa = 0 -\n non-flat ΛCDM: w0 = -1, wa = 0 - wCDM: wa = 0 - w0waCDM: general case\n\n All computations use Nx tensor ops, making them natively differentiable\n through Rune's autodiff. GL quadrature is vectorized as tensor operations. *)\n\nlet f64 = Nx.float64\nlet c_km_s = Nx.scalar f64 299792.458\nlet _mpc_m = 3.085_677_581_491_367_3e22\n\ntype params = {\n h0 : Nx.float64_t;\n omega_m : Nx.float64_t;\n omega_l : Nx.float64_t;\n omega_k : Nx.float64_t;\n w0 : Nx.float64_t;\n wa : Nx.float64_t;\n omega_b : Nx.float64_t option;\n n_s : Nx.float64_t option;\n sigma8 : Nx.float64_t option;\n}\n\nlet err_missing name =\n invalid_arg\n (\"Cosmo: \" ^ name ^ \" not set (use Cosmo.set or a preset like planck18)\")\n\n(* --- Constructors --- *)\n\nlet flat_lcdm ~h0 ~omega_m =\n if h0 <= 0.0 then invalid_arg \"Cosmo.flat_lcdm: h0 must be positive\";\n if omega_m < 0.0 then\n invalid_arg \"Cosmo.flat_lcdm: omega_m must be non-negative\";\n {\n h0 = Nx.scalar f64 h0;\n omega_m = Nx.scalar f64 omega_m;\n omega_l = Nx.scalar f64 (1.0 -. omega_m);\n omega_k = Nx.scalar f64 0.0;\n w0 = Nx.scalar f64 (-1.0);\n wa = Nx.scalar f64 0.0;\n omega_b = None;\n n_s = None;\n sigma8 = None;\n }\n\nlet lcdm ~h0 ~omega_m ~omega_l =\n if h0 <= 0.0 then invalid_arg \"Cosmo.lcdm: h0 must be positive\";\n {\n h0 = Nx.scalar f64 h0;\n omega_m = Nx.scalar f64 omega_m;\n omega_l = Nx.scalar f64 omega_l;\n omega_k = Nx.scalar f64 (1.0 -. omega_m -. omega_l);\n w0 = Nx.scalar f64 (-1.0);\n wa = Nx.scalar f64 0.0;\n omega_b = None;\n n_s = None;\n sigma8 = None;\n }\n\nlet wcdm ~h0 ~omega_m ?omega_l ~w0 () =\n if h0 <= 0.0 then invalid_arg \"Cosmo.wcdm: h0 must be positive\";\n let omega_l = match omega_l with Some v -> v | None -> 1.0 -. omega_m in\n {\n h0 = Nx.scalar f64 h0;\n omega_m = Nx.scalar f64 omega_m;\n omega_l = Nx.scalar f64 omega_l;\n omega_k = Nx.scalar f64 (1.0 -. omega_m -. omega_l);\n w0 = Nx.scalar f64 w0;\n wa = Nx.scalar f64 0.0;\n omega_b = None;\n n_s = None;\n sigma8 = None;\n }\n\nlet w0wacdm ~h0 ~omega_m ?omega_l ~w0 ~wa () =\n if h0 <= 0.0 then invalid_arg \"Cosmo.w0wacdm: h0 must be positive\";\n let omega_l = match omega_l with Some v -> v | None -> 1.0 -. omega_m in\n {\n h0 = Nx.scalar f64 h0;\n omega_m = Nx.scalar f64 omega_m;\n omega_l = Nx.scalar f64 omega_l;\n omega_k = Nx.scalar f64 (1.0 -. omega_m -. omega_l);\n w0 = Nx.scalar f64 w0;\n wa = Nx.scalar f64 wa;\n omega_b = None;\n n_s = None;\n sigma8 = None;\n }\n\n(* Tensor constructors for differentiable construction *)\n\nlet create_flat_lcdm ~h0 ~omega_m =\n {\n h0;\n omega_m;\n omega_l = Nx.sub (Nx.scalar f64 1.0) omega_m;\n omega_k = Nx.scalar f64 0.0;\n w0 = Nx.scalar f64 (-1.0);\n wa = Nx.scalar f64 0.0;\n omega_b = None;\n n_s = None;\n sigma8 = None;\n }\n\nlet create_lcdm ~h0 ~omega_m ~omega_l =\n {\n h0;\n omega_m;\n omega_l;\n omega_k = Nx.sub (Nx.scalar f64 1.0) (Nx.add omega_m omega_l);\n w0 = Nx.scalar f64 (-1.0);\n wa = Nx.scalar f64 0.0;\n omega_b = None;\n n_s = None;\n sigma8 = None;\n }\n\nlet create_wcdm ~h0 ~omega_m ?omega_l ~w0 () =\n let omega_l =\n match omega_l with\n | Some v -> v\n | None -> Nx.sub (Nx.scalar f64 1.0) omega_m\n in\n {\n h0;\n omega_m;\n omega_l;\n omega_k = Nx.sub (Nx.scalar f64 1.0) (Nx.add omega_m omega_l);\n w0;\n wa = Nx.scalar f64 0.0;\n omega_b = None;\n n_s = None;\n sigma8 = None;\n }\n\nlet create_w0wacdm ~h0 ~omega_m ?omega_l ~w0 ~wa () =\n let omega_l =\n match omega_l with\n | Some v -> v\n | None -> Nx.sub (Nx.scalar f64 1.0) omega_m\n in\n {\n h0;\n omega_m;\n omega_l;\n omega_k = Nx.sub (Nx.scalar f64 1.0) (Nx.add omega_m omega_l);\n w0;\n wa;\n omega_b = None;\n n_s = None;\n sigma8 = None;\n }\n\n(* Accessors *)\n\nlet h0 p = p.h0\nlet omega_m p = p.omega_m\nlet omega_l p = p.omega_l\nlet omega_k p = p.omega_k\nlet w0 p = p.w0\nlet wa p = p.wa\n\nlet omega_b p =\n match p.omega_b with Some v -> v | None -> err_missing \"omega_b\"\n\nlet n_s p = match p.n_s with Some v -> v | None -> err_missing \"n_s\"\nlet sigma8 p = match p.sigma8 with Some v -> v | None -> err_missing \"sigma8\"\n\nlet set ?omega_b ?n_s ?sigma8 p =\n let omega_b =\n match omega_b with Some v -> Some (Nx.scalar f64 v) | None -> p.omega_b\n in\n let n_s = match n_s with Some v -> Some (Nx.scalar f64 v) | None -> p.n_s in\n let sigma8 =\n match sigma8 with Some v -> Some (Nx.scalar f64 v) | None -> p.sigma8\n in\n { p with omega_b; n_s; sigma8 }\n\nlet set_t ?h0 ?omega_m ?omega_l ?omega_b ?n_s ?sigma8 p =\n let h0 = match h0 with Some v -> v | None -> p.h0 in\n let omega_m = match omega_m with Some v -> v | None -> p.omega_m in\n let omega_l = match omega_l with Some v -> v | None -> p.omega_l in\n let omega_k = Nx.sub (Nx.scalar f64 1.0) (Nx.add omega_m omega_l) in\n let omega_b = match omega_b with Some v -> Some v | None -> p.omega_b in\n let n_s = match n_s with Some v -> Some v | None -> p.n_s in\n let sigma8 = match sigma8 with Some v -> Some v | None -> p.sigma8 in\n { p with h0; omega_m; omega_l; omega_k; omega_b; n_s; sigma8 }\n\n(* Presets *)\n\nlet default = flat_lcdm ~h0:70.0 ~omega_m:0.3\n\nlet planck18 =\n flat_lcdm ~h0:67.66 ~omega_m:0.3111\n |> set ~omega_b:0.0490 ~n_s:0.9665 ~sigma8:0.8102\n\nlet planck15 =\n flat_lcdm ~h0:67.74 ~omega_m:0.3075\n |> set ~omega_b:0.0486 ~n_s:0.9667 ~sigma8:0.8159\n\nlet wmap9 =\n flat_lcdm ~h0:69.32 ~omega_m:0.2865\n |> set ~omega_b:0.0463 ~n_s:0.9608 ~sigma8:0.820\n\n(* --- E(z) computation ---\n\n E(z) = H(z)/H0 = sqrt(Ω_m(1+z)³ + Ω_k(1+z)² + Ω_de(z))\n\n where Ω_de(z) = Ω_Λ * (1+z)^(3(1+w0+wa)) * exp(-3*wa*z/(1+z))\n\n For ΛCDM (w0=-1, wa=0): Ω_de(z) = Ω_Λ (constant) For wCDM (wa=0): Ω_de(z) =\n Ω_Λ * (1+z)^(3(1+w0)) *)\n\nlet e_of p z =\n let one_plus_z = Nx.add_s z 1.0 in\n let cubed = Nx.mul one_plus_z (Nx.mul one_plus_z one_plus_z) in\n let matter = Nx.mul p.omega_m cubed in\n let curvature = Nx.mul p.omega_k (Nx.mul one_plus_z one_plus_z) in\n (* Dark energy: Ω_Λ * (1+z)^(3(1+w0+wa)) * exp(-3*wa*z/(1+z)) *)\n let w_eff = Nx.add_s (Nx.add p.w0 p.wa) 1.0 in\n let de_power = Nx.pow one_plus_z (Nx.mul_s w_eff 3.0) in\n let wa_arg = Nx.mul (Nx.mul_s p.wa (-3.0)) (Nx.div z one_plus_z) in\n let de = Nx.mul p.omega_l (Nx.mul de_power (Nx.exp wa_arg)) in\n Nx.sqrt (Nx.add matter (Nx.add curvature de))\n\n(* 16-point Gauss-Legendre nodes and weights on [-1, 1] as Nx tensors *)\nlet gl_nodes =\n Nx.create f64 [| 16 |]\n [|\n -0.9894009349916499;\n -0.9445750230732326;\n -0.8656312023878318;\n -0.7554044083550030;\n -0.6178762444026438;\n -0.4580167776572274;\n -0.2816035507792589;\n -0.0950125098376374;\n 0.0950125098376374;\n 0.2816035507792589;\n 0.4580167776572274;\n 0.6178762444026438;\n 0.7554044083550030;\n 0.8656312023878318;\n 0.9445750230732326;\n 0.9894009349916499;\n |]\n\nlet gl_weights =\n Nx.create f64 [| 16 |]\n [|\n 0.0271524594117541;\n 0.0622535239386479;\n 0.0951585116824928;\n 0.1246289712555339;\n 0.1495959888165767;\n 0.1691565193950025;\n 0.1826034150449236;\n 0.1894506104550685;\n 0.1894506104550685;\n 0.1826034150449236;\n 0.1691565193950025;\n 0.1495959888165767;\n 0.1246289712555339;\n 0.0951585116824928;\n 0.0622535239386479;\n 0.0271524594117541;\n |]\n\n(* GL quadrature in scale-factor space.\n\n All cosmological integrals ∫₀ᶻ g(z') dz' are evaluated via the substitution a\n = 1/(1+z), which maps [0, z] → [1/(1+z), 1]. This bounded range is\n well-resolved by 16-point GL even at z = 1089 (CMB). Direct quadrature over\n [0, z] in redshift space under-resolves the integrand at large z. *)\nlet gl_quad_a p z f =\n let a_lo = Nx.recip (Nx.add_s z 1.0) in\n let one = Nx.scalar f64 1.0 in\n let half = Nx.div_s (Nx.sub one a_lo) 2.0 in\n let mid = Nx.div_s (Nx.add one a_lo) 2.0 in\n let a = Nx.add (Nx.mul half gl_nodes) mid in\n let e_z = e_of p (Nx.sub_s (Nx.recip a) 1.0) in\n Nx.mul half (Nx.sum (Nx.mul (f a e_z) gl_weights))\n\n(* ∫₀ᶻ dz'/E(z') = ∫_{a_lo}^1 da/(a² E(a)) *)\nlet integrate_inv_ez p z =\n gl_quad_a p z (fun a e -> Nx.recip (Nx.mul (Nx.mul a a) e))\n\n(* ∫₀ᶻ dz'/((1+z') E(z')) = ∫_{a_lo}^1 da/(a E(a)) *)\nlet integrate_inv_z1_ez p z = gl_quad_a p z (fun a e -> Nx.recip (Nx.mul a e))\n\n(* --- Derived quantities --- *)\n\nlet hubble ?(p = default) z = Nx.mul p.h0 (e_of p z)\n\nlet critical_density ?(p = default) z =\n let h_z = hubble ~p z in\n let h_si = Nx.div_s (Nx.mul_s h_z 1e3) _mpc_m in\n Nx.div_s (Nx.mul_s (Nx.mul h_si h_si) 3.0) (8.0 *. Float.pi *. 6.674_30e-11)\n\n(* --- Distances ---\n\n Line-of-sight comoving distance: χ = d_H ∫₀ᶻ dz'/E(z')\n\n Transverse comoving distance (curvature-corrected): - Ω_k > 0 (open): d_M =\n d_H/√Ω_k · sinh(√Ω_k · χ/d_H) - Ω_k = 0 (flat): d_M = χ - Ω_k < 0 (closed):\n d_M = d_H/√|Ω_k| · sin(√|Ω_k| · χ/d_H) *)\n\nlet comoving_distance_mpc p z =\n let d_h = Nx.div c_km_s p.h0 in\n Nx.mul d_h (integrate_inv_ez p z)\n\nlet transverse_comoving_mpc p z =\n let d_h = Nx.div c_km_s p.h0 in\n let chi = Nx.mul d_h (integrate_inv_ez p z) in\n let ok_f = Nx.item [] p.omega_k in\n if Float.abs ok_f < 1e-10 then chi (* flat *)\n else\n let sqrt_ok = Nx.sqrt (Nx.abs p.omega_k) in\n let arg = Nx.div (Nx.mul sqrt_ok chi) d_h in\n if ok_f > 0.0 then Nx.div (Nx.mul d_h (Nx.sinh arg)) sqrt_ok\n else Nx.div (Nx.mul d_h (Nx.sin arg)) sqrt_ok\n\nlet comoving_distance ?(p = default) z =\n Unit.Length.of_tensor (Nx.mul_s (comoving_distance_mpc p z) _mpc_m)\n\nlet luminosity_distance ?(p = default) z =\n let dm_mpc = transverse_comoving_mpc p z in\n Unit.Length.of_tensor (Nx.mul_s (Nx.mul (Nx.add_s z 1.0) dm_mpc) _mpc_m)\n\nlet angular_diameter_distance ?(p = default) z =\n let dm_mpc = transverse_comoving_mpc p z in\n Unit.Length.of_tensor (Nx.mul_s (Nx.div dm_mpc (Nx.add_s z 1.0)) _mpc_m)\n\nlet distance_modulus ?(p = default) z =\n let dl_mpc = Nx.mul (Nx.add_s z 1.0) (transverse_comoving_mpc p z) in\n (* mu = 5 * log10(dL_Mpc) + 25 = 5/ln10 * ln(dL_Mpc) + 25 *)\n let five_over_ln10 = 5.0 /. Float.log 10.0 in\n Nx.add_s (Nx.mul_s (Nx.log dl_mpc) five_over_ln10) 25.0\n\n(* --- Angular scale --- *)\n\nlet angular_size ?(p = default) ~z phys =\n let da = angular_diameter_distance ~p z in\n Unit.Angle.of_tensor\n (Nx.div (Unit.Length.to_tensor phys) (Unit.Length.to_tensor da))\n\nlet physical_size ?(p = default) ~z ang =\n let da = angular_diameter_distance ~p z in\n Unit.Length.of_tensor\n (Nx.mul (Unit.Angle.to_tensor ang) (Unit.Length.to_tensor da))\n\n(* --- Cosmic times --- *)\n\n(* 1/H0 in seconds: (km/s/Mpc)^{-1} = Mpc/km · s *)\nlet _hubble_time_s p = Nx.mul_s (Nx.recip p.h0) 3.0856776e19\n\nlet lookback_time ?(p = default) z =\n Unit.Time.of_tensor (Nx.mul (_hubble_time_s p) (integrate_inv_z1_ez p z))\n\nlet age ?(p = default) z =\n (* age(z) = t_H ∫₀^{1/(1+z)} da/(a E(a)). We reuse gl_quad_a with an upper\n limit at z_max=1000 (≈ a_lo → 0) for the total integral, then subtract the\n lookback from 0 to z. *)\n let t_h_s = _hubble_time_s p in\n let total = integrate_inv_z1_ez p (Nx.scalar f64 1000.0) in\n let lb = integrate_inv_z1_ez p z in\n Unit.Time.of_tensor (Nx.mul t_h_s (Nx.sub total lb))\n\n(* --- z_at_value: inverse lookup via Brent's method ---\n\n Given a monotonic cosmological function f and a target value, find the\n redshift z such that f(z) ≈ target. Not differentiable. *)\n\nlet z_at_value ?(p = default) ?(zmin = 1e-8) ?(zmax = 1000.0) ?(xtol = 1e-8) f\n target =\n let target_v = Nx.item [] target in\n let eval z = Nx.item [] (f ~p (Nx.scalar f64 z)) -. target_v in\n (* Brent's method *)\n let a = ref zmin and b = ref zmax in\n let fa = ref (eval !a) and fb = ref (eval !b) in\n if !fa *. !fb > 0.0 then\n invalid_arg \"Cosmo.z_at_value: target outside [f(zmin), f(zmax)]\";\n if Float.abs !fa < Float.abs !fb then begin\n let tmp = !a in\n a := !b;\n b := tmp;\n let tmp = !fa in\n fa := !fb;\n fb := tmp\n end;\n let c = ref !a and fc = ref !fa in\n let d = ref (!b -. !a) in\n let mflag = ref true in\n let max_iter = 100 in\n let i = ref 0 in\n while Float.abs !fb > xtol && !i < max_iter do\n let s =\n if Float.abs (!fa -. !fc) > 1e-30 && Float.abs (!fb -. !fc) > 1e-30 then\n (* Inverse quadratic interpolation *)\n let s1 = !a *. !fb *. !fc /. ((!fa -. !fb) *. (!fa -. !fc)) in\n let s2 = !b *. !fa *. !fc /. ((!fb -. !fa) *. (!fb -. !fc)) in\n let s3 = !c *. !fa *. !fb /. ((!fc -. !fa) *. (!fc -. !fb)) in\n s1 +. s2 +. s3\n else\n (* Secant method *)\n !b -. (!fb *. (!b -. !a) /. (!fb -. !fa))\n in\n let cond1 =\n let lo = ((3.0 *. !a) +. !b) /. 4.0 in\n not (if lo < !b then lo <= s && s <= !b else !b <= s && s <= lo)\n in\n let cond2 = !mflag && Float.abs (s -. !b) >= Float.abs (!b -. !c) /. 2.0 in\n let cond3 =\n (not !mflag) && Float.abs (s -. !b) >= Float.abs (!c -. !d) /. 2.0\n in\n let cond4 = !mflag && Float.abs (!b -. !c) < xtol in\n let cond5 = (not !mflag) && Float.abs (!c -. !d) < xtol in\n let s =\n if cond1 || cond2 || cond3 || cond4 || cond5 then begin\n mflag := true;\n (!a +. !b) /. 2.0\n end\n else begin\n mflag := false;\n s\n end\n in\n let fs = eval s in\n d := !c;\n c := !b;\n fc := !fb;\n if !fa *. fs < 0.0 then begin\n b := s;\n fb := fs\n end\n else begin\n a := s;\n fa := fs\n end;\n if Float.abs !fa < Float.abs !fb then begin\n let tmp = !a in\n a := !b;\n b := tmp;\n let tmp = !fa in\n fa := !fb;\n fb := tmp\n end;\n incr i\n done;\n Nx.scalar f64 !b\n\n(* Growth factor and growth rate *)\n\n(* E(a) from scale factor: a = 1/(1+z), so z = 1/a - 1 *)\nlet e_at_a p a = e_of p (Nx.sub_s (Nx.recip a) 1.0)\n\n(* GL quadrature of f(a') from 0 to a. Transforms [-1,1] to [0,a]. *)\nlet gl_integrate_a p a f =\n let half = Nx.div_s a 2.0 in\n let a_prime = Nx.add (Nx.mul half gl_nodes) half in\n let e_a = e_at_a p a_prime in\n Nx.mul half (Nx.sum (Nx.mul (f a_prime e_a) gl_weights))\n\n(* Growth integral: J(a) = ∫₀ᵃ da' / (a'³ E³(a')) Integrand at a'→0:\n ~a'^(3/2)/Ω_m^(3/2) → 0, so well-behaved. *)\nlet growth_integral p a =\n gl_integrate_a p a (fun a_prime e_a ->\n let a3 = Nx.mul a_prime (Nx.mul a_prime a_prime) in\n let e3 = Nx.mul e_a (Nx.mul e_a e_a) in\n Nx.recip (Nx.mul a3 e3))\n\n(* Unnormalized growth factor: D(a) ∝ E(a) × J(a) *)\nlet growth_unnorm p a = Nx.mul (e_at_a p a) (growth_integral p a)\n\nlet growth_factor ?(p = default) z =\n let a = Nx.recip (Nx.add_s z 1.0) in\n let d_a = growth_unnorm p a in\n let d_1 = growth_unnorm p (Nx.scalar f64 1.0) in\n Nx.div d_a d_1\n\n(* Growth rate: f(a) = dlnD/dlna D(a) = E(a) J(a) / const, so f = dlnE/dlna +\n (dJ/dlna) / J = dlnE/dlna + 1 / (a² E³(a) J(a))\n\n dlnE/dlna = a/(2E²) dE²/da dE²/da = -3Ωm a⁻⁴ - 2Ωk a⁻³ + ΩΛ exp(f_de)\n (-3(1+w0+wa)/a + 3wa) *)\nlet growth_rate ?(p = default) z =\n let a = Nx.recip (Nx.add_s z 1.0) in\n let e_a = e_at_a p a in\n let e2 = Nx.mul e_a e_a in\n let j_a = growth_integral p a in\n (* dE²/da *)\n let a2 = Nx.mul a a in\n let a3 = Nx.mul a2 a in\n let a4 = Nx.mul a3 a in\n let dm = Nx.mul_s (Nx.div p.omega_m a4) (-3.0) in\n let dk = Nx.mul_s (Nx.div p.omega_k a3) (-2.0) in\n (* Dark energy contribution: need f_de(a) and f_de'(a) *)\n let f_de =\n Nx.add\n (Nx.mul (Nx.mul_s (Nx.add_s (Nx.add p.w0 p.wa) 1.0) (-3.0)) (Nx.log a))\n (Nx.mul p.wa (Nx.mul_s (Nx.sub_s a 1.0) 3.0))\n in\n let f_de_prime =\n Nx.add\n (Nx.div (Nx.mul_s (Nx.add_s (Nx.add p.w0 p.wa) 1.0) (-3.0)) a)\n (Nx.mul_s p.wa 3.0)\n in\n let dde = Nx.mul (Nx.mul p.omega_l (Nx.exp f_de)) f_de_prime in\n let de2_da = Nx.add dm (Nx.add dk dde) in\n (* dlnE/dlna = a/(2E²) × dE²/da *)\n let dln_e = Nx.div (Nx.mul a de2_da) (Nx.mul_s e2 2.0) in\n (* 1/(a² E³ J) *)\n let e3 = Nx.mul e_a e2 in\n let term2 = Nx.recip (Nx.mul a2 (Nx.mul e3 j_a)) in\n Nx.add dln_e term2\n\n(* Eisenstein-Hu transfer function (1998) *)\n\nlet t_cmb = 2.7255\n\n(* Eisenstein & Hu (1998) transfer function with baryon oscillations. Scalar\n cosmological quantities are computed in float arithmetic (the transfer\n function is a fitting formula. The wavenumber k may be a tensor of arbitrary\n shape; the result has the same shape. Differentiable through cosmological\n parameters via Rune. *)\nlet eisenstein_hu p k =\n let s = Nx.scalar f64 in\n let om = p.omega_m in\n let ob = omega_b p in\n let h = Nx.div_s p.h0 100.0 in\n let h2 = Nx.mul h h in\n let w_m = Nx.mul om h2 in\n let w_b = Nx.mul ob h2 in\n let fb = Nx.div ob om in\n let fc = Nx.sub (s 1.0) fb in\n let t27sq = (t_cmb /. 2.7) ** 2.0 in\n let t27_4 = t27sq *. t27sq in\n (* Eq. 2,3: equality epoch *)\n let z_eq = Nx.div_s (Nx.mul_s w_m 2.50e4) t27_4 in\n let k_eq = Nx.div (Nx.div_s (Nx.mul_s w_m 7.46e-2) t27sq) h in\n (* Eq. 4: drag epoch *)\n let b1 =\n Nx.mul\n (Nx.pow w_m (s (-0.419)))\n (Nx.add_s (Nx.mul_s (Nx.pow w_m (s 0.674)) 0.607) 1.0)\n |> fun x -> Nx.mul_s x 0.313\n in\n let b2 = Nx.mul_s (Nx.pow w_m (s 0.223)) 0.238 in\n let z_d =\n Nx.mul\n (Nx.div\n (Nx.mul_s (Nx.pow w_m (s 0.251)) 1291.0)\n (Nx.add_s (Nx.mul_s (Nx.pow w_m (s 0.828)) 0.659) 1.0))\n (Nx.add_s (Nx.mul b1 (Nx.pow w_b b2)) 1.0)\n in\n (* Eq. 5: baryon/photon momentum ratios *)\n let r_d = Nx.mul (Nx.div_s (Nx.mul_s w_b 31.5) t27_4) (Nx.div (s 1e3) z_d) in\n let r_eq =\n Nx.mul (Nx.div_s (Nx.mul_s w_b 31.5) t27_4) (Nx.div (s 1e3) z_eq)\n in\n (* Eq. 6: sound horizon *)\n let sh_d =\n Nx.mul\n (Nx.mul\n (Nx.div (s 2.0) (Nx.mul_s k_eq 3.0))\n (Nx.sqrt (Nx.div (s 6.0) r_eq)))\n (Nx.log\n (Nx.div\n (Nx.add (Nx.sqrt (Nx.add_s r_d 1.0)) (Nx.sqrt (Nx.add r_eq r_d)))\n (Nx.add_s (Nx.sqrt r_eq) 1.0)))\n in\n (* Eq. 7: Silk damping *)\n let k_silk =\n Nx.div\n (Nx.mul\n (Nx.mul (Nx.mul_s (Nx.pow w_b (s 0.52)) 1.6) (Nx.pow w_m (s 0.73)))\n (Nx.add_s (Nx.pow (Nx.mul_s w_m 10.4) (s (-0.95))) 1.0))\n h\n in\n (* CDM transfer function (Eqs. 11, 12, 17, 18) *)\n let a1 =\n Nx.mul\n (Nx.pow (Nx.mul_s w_m 46.9) (s 0.670))\n (Nx.add_s (Nx.pow (Nx.mul_s w_m 32.1) (s (-0.532))) 1.0)\n in\n let a2 =\n Nx.mul\n (Nx.pow (Nx.mul_s w_m 12.0) (s 0.424))\n (Nx.add_s (Nx.pow (Nx.mul_s w_m 45.0) (s (-0.582))) 1.0)\n in\n let alpha_c =\n Nx.mul\n (Nx.pow a1 (Nx.neg fb))\n (Nx.pow a2 (Nx.neg (Nx.mul fb (Nx.mul fb fb))))\n in\n let b1c =\n Nx.div (s 0.944) (Nx.add_s (Nx.pow (Nx.mul_s w_m 458.0) (s (-0.708))) 1.0)\n in\n let b2c = Nx.pow (Nx.mul_s w_m 0.395) (s (-0.0266)) in\n let beta_c =\n Nx.recip (Nx.add_s (Nx.mul b1c (Nx.sub (Nx.pow fc b2c) (s 1.0))) 1.0)\n in\n (* T_tilde: Eq. 10, 19. Operates on k tensor. alpha, beta are scalar\n tensors. *)\n let t_tilde k1 alpha beta =\n let q = Nx.div k1 (Nx.mul_s k_eq 13.41) in\n let l = Nx.log (Nx.add_s (Nx.mul q (Nx.mul_s beta 1.8)) (Float.exp 1.0)) in\n let c =\n Nx.add\n (Nx.div (s 386.0) (Nx.add_s (Nx.mul_s (Nx.pow q (s 1.08)) 69.9) 1.0))\n (Nx.div (s 14.2) alpha)\n in\n Nx.div l (Nx.add l (Nx.mul c (Nx.mul q q)))\n in\n let ksh = Nx.mul k sh_d in\n (* Eq. 17, 18 *)\n let f_ =\n let x = Nx.div_s ksh 5.4 in\n let x2 = Nx.mul x x in\n Nx.recip (Nx.add_s (Nx.mul x2 x2) 1.0)\n in\n let tc =\n Nx.add\n (Nx.mul f_ (t_tilde k (s 1.0) beta_c))\n (Nx.mul (Nx.sub (s 1.0) f_) (t_tilde k alpha_c beta_c))\n in\n (* Baryon transfer function (Eqs. 14, 19, 21) *)\n let y = Nx.div (Nx.add_s z_eq 1.0) (Nx.add_s z_d 1.0) in\n let x_ = Nx.sqrt (Nx.add_s y 1.0) in\n let g_eh =\n Nx.mul y\n (Nx.add (Nx.mul_s x_ (-6.0))\n (Nx.mul\n (Nx.add_s (Nx.mul_s y 3.0) 2.0)\n (Nx.log (Nx.div (Nx.add_s x_ 1.0) (Nx.sub_s x_ 1.0)))))\n in\n let alpha_b =\n Nx.mul_s\n (Nx.mul (Nx.mul k_eq sh_d)\n (Nx.mul (Nx.pow (Nx.add_s r_d 1.0) (s (-0.75))) g_eh))\n 2.07\n in\n let beta_node = Nx.mul_s (Nx.pow w_m (s 0.435)) 8.41 in\n let beta_b =\n Nx.add (Nx.add_s fb 0.5)\n (Nx.mul\n (Nx.sub_s (Nx.mul_s fb 2.0) 3.0)\n (Nx.neg\n (Nx.sqrt\n (Nx.add_s (Nx.mul (Nx.mul_s w_m 17.2) (Nx.mul_s w_m 17.2)) 1.0))))\n in\n (* Eq. 22: tilde_s per-k *)\n let tilde_s =\n let bns = Nx.div beta_node ksh in\n let bns3 = Nx.mul bns (Nx.mul bns bns) in\n Nx.div sh_d (Nx.pow (Nx.add_s bns3 1.0) (s (1.0 /. 3.0)))\n in\n let tb =\n let term1 =\n Nx.div\n (t_tilde k (s 1.0) (s 1.0))\n (Nx.add_s\n (let x = Nx.div_s ksh 5.2 in\n Nx.mul x x)\n 1.0)\n in\n let bbks = Nx.div beta_b ksh in\n let bbks3 = Nx.mul bbks (Nx.mul bbks bbks) in\n let term2 =\n Nx.mul\n (Nx.div alpha_b (Nx.add_s bbks3 1.0))\n (Nx.exp (Nx.neg (Nx.pow (Nx.div k k_silk) (s 1.4))))\n in\n let sinc_arg = Nx.mul k tilde_s in\n Nx.mul (Nx.add term1 term2) (Nx.div (Nx.sin sinc_arg) sinc_arg)\n in\n (* Total: fb * Tb + fc * Tc *)\n Nx.add (Nx.mul tb fb) (Nx.mul tc fc)\n\n(* Matter power spectrum *)\n\n(* Simpson's rule integration on a uniform grid of n+1 points from a to b. n\n must be even. f is evaluated at each grid point, returns [n+1] tensor. *)\nlet simps_integrate f a b n =\n let h = (b -. a) /. Float.of_int n in\n let xs =\n Nx.create f64\n [| n + 1 |]\n (Array.init (n + 1) (fun i -> a +. (Float.of_int i *. h)))\n in\n let ys = f xs in\n (* Simpson weights: 1, 4, 2, 4, 2, ..., 4, 1 *)\n let w =\n Array.init (n + 1) (fun i ->\n if i = 0 || i = n then 1.0 else if i mod 2 = 1 then 4.0 else 2.0)\n in\n let weights = Nx.create f64 [| n + 1 |] w in\n Nx.mul_s (Nx.sum (Nx.mul ys weights)) (h /. 3.0)\n\n(* σ²(R) = 1/(2π²) ∫ k³ P_unnorm(k) W²(kR) d(ln k) where P_unnorm = k^n_s ×\n T²(k) and W is the top-hat window. Integration in ln(k) space: the integrand\n is k³ P W² (the dk/k from d(ln k) cancels one power of k, giving k² P W² dk\n equivalent). *)\nlet sigma_sq p r =\n let ns = n_s p in\n simps_integrate\n (fun lnk ->\n let k = Nx.exp lnk in\n let x = Nx.mul_s k r in\n (* Top-hat window: W(x) = 3(sin x - x cos x)/x³ *)\n let x2 = Nx.mul x x in\n let x3 = Nx.mul x2 x in\n let w =\n Nx.div (Nx.mul_s (Nx.sub (Nx.sin x) (Nx.mul x (Nx.cos x))) 3.0) x3\n in\n let t = eisenstein_hu p k in\n let pk = Nx.mul (Nx.pow k ns) (Nx.mul t t) in\n let k3 = Nx.mul k (Nx.mul k k) in\n Nx.mul k3 (Nx.mul (Nx.mul w w) pk))\n (Float.log 1e-4) (Float.log 1e4) 512\n |> fun integral -> Nx.div_s integral (2.0 *. Float.pi *. Float.pi)\n\nlet linear_power ?(p = default) k z =\n let s8 = sigma8 p in\n let g = growth_factor ~p z in\n let t = eisenstein_hu p k in\n let ns = n_s p in\n let pk_unnorm = Nx.mul (Nx.pow k ns) (Nx.mul t t) in\n (* Normalization: A = σ8² / σ²_unnorm(R=8) *)\n let s2 = sigma_sq p 8.0 in\n let norm = Nx.div (Nx.mul s8 s8) s2 in\n Nx.mul norm (Nx.mul pk_unnorm (Nx.mul g g))\n\n(* Halofit (Takahashi et al. 2012) *)\n\n(* Ω_m(a) = Ω_m a⁻³ / E²(a) *)\nlet omega_m_a p a =\n let e2 =\n let e = e_at_a p a in\n Nx.mul e e\n in\n let a3 = Nx.mul a (Nx.mul a a) in\n Nx.div (Nx.div p.omega_m a3) e2\n\n(* Ω_de(a) = Ω_Λ exp(f_de(a)) / E²(a) *)\nlet omega_de_a p a =\n let e2 =\n let e = e_at_a p a in\n Nx.mul e e\n in\n let f_de =\n Nx.add\n (Nx.mul (Nx.mul_s (Nx.add_s (Nx.add p.w0 p.wa) 1.0) (-3.0)) (Nx.log a))\n (Nx.mul p.wa (Nx.mul_s (Nx.sub_s a 1.0) 3.0))\n in\n Nx.div (Nx.mul p.omega_l (Nx.exp f_de)) e2\n\n(* w(a) = w0 + wa(1-a) *)\nlet w_of p a = Nx.add p.w0 (Nx.mul p.wa (Nx.sub (Nx.scalar f64 1.0) a))\n\n(* σ²(R, z) using linear P(k) at z=0, scaled by D²(z)/D²(0)=D²(z). For Halofit\n we need σ(R) at various R to find k_nl, plus derivatives. *)\nlet sigma_sq_at_z p r z =\n let g = growth_factor ~p z in\n Nx.mul (sigma_sq p r) (Nx.mul g g)\n\n(* Find k_nl where σ(1/k_nl, z) = 1, plus n_eff and C at the nonlinear scale. We\n compute σ²(R) on a grid, interpolate to find R_nl, then compute spectral\n index and curvature from Gaussian-filtered integrals. *)\nlet halofit_params p z =\n let g = growth_factor ~p z in\n let g2 = Nx.mul g g in\n let ns = n_s p in\n let s8 = sigma8 p in\n let s2_8 = sigma_sq p 8.0 in\n let pknorm = Nx.div (Nx.mul s8 s8) s2_8 in\n let n_r = 256 in\n let logr =\n Nx.create f64 [| n_r |]\n (Array.init n_r (fun i ->\n Float.log 1e-4\n +. Float.of_int i\n *. (Float.log 1e1 -. Float.log 1e-4)\n /. Float.of_int (n_r - 1)))\n in\n (* Compute σ²(R) for each R using Gaussian filter exp(-(kR)²) *)\n let n_k = 512 in\n let lnk_min = Float.log 1e-4 in\n let lnk_max = Float.log 1e4 in\n let dlnk = (lnk_max -. lnk_min) /. Float.of_int (n_k - 1) in\n let lnk =\n Nx.create f64 [| n_k |]\n (Array.init n_k (fun i -> lnk_min +. (Float.of_int i *. dlnk)))\n in\n let k = Nx.exp lnk in\n let t = eisenstein_hu p k in\n let pk_base = Nx.mul pknorm (Nx.mul (Nx.pow k ns) (Nx.mul t t)) in\n let pk_at_z = Nx.mul pk_base g2 in\n (* k³ P(k) / (2π²) *)\n let k3pk =\n Nx.div\n (Nx.mul (Nx.mul k (Nx.mul k k)) pk_at_z)\n (Nx.scalar f64 (2.0 *. Float.pi *. Float.pi))\n in\n (* Trapezoidal weights [n_k] *)\n let trap_w =\n Nx.create f64 [| n_k |]\n (Array.init n_k (fun j -> if j = 0 || j = n_k - 1 then 0.5 else 1.0))\n in\n (* Float-level σ²(R) grid for root-finding *)\n let sigma2_arr = Array.make n_r 0.0 in\n for i = 0 to n_r - 1 do\n let r = Float.exp (Nx.item [ i ] logr) in\n let kr = Nx.mul_s k r in\n let y2 = Nx.mul kr kr in\n let gauss = Nx.exp (Nx.neg y2) in\n let integrand = Nx.mul k3pk gauss in\n sigma2_arr.(i) <-\n Nx.item [] (Nx.mul_s (Nx.sum (Nx.mul trap_w integrand)) dlnk)\n done;\n (* Find R_nl where σ² = 1 by linear interpolation in log space *)\n let r_nl = ref (Float.exp (Nx.item [ 0 ] logr)) in\n (let found = ref false in\n for i = 0 to n_r - 2 do\n if (not !found) && sigma2_arr.(i) >= 1.0 && sigma2_arr.(i + 1) <= 1.0 then begin\n let ls0 = Float.log sigma2_arr.(i) in\n let ls1 = Float.log sigma2_arr.(i + 1) in\n let lr0 = Nx.item [ i ] logr in\n let lr1 = Nx.item [ i + 1 ] logr in\n let frac = (0.0 -. ls0) /. (ls1 -. ls0) in\n r_nl := Float.exp (lr0 +. (frac *. (lr1 -. lr0)));\n found := true\n end\n done);\n let r_nl_f = !r_nl in\n (* Differentiable Newton refinement for R_nl. Compute σ² at the float root,\n then one Newton step: R' = R + R*(σ²-1)/dn where dσ²/dR = -dn/R.\n Numerically R' ≈ R, but the gradient dR'/dp = -(∂σ²/∂p)/(∂σ²/∂R) is exact\n via the implicit function theorem. *)\n let kr0 = Nx.mul_s k r_nl_f in\n let y2_0 = Nx.mul kr0 kr0 in\n let gauss0 = Nx.exp (Nx.neg y2_0) in\n let integrand0 = Nx.mul k3pk gauss0 in\n let trap_sum f = Nx.mul_s (Nx.sum (Nx.mul trap_w f)) dlnk in\n let s2_0 = trap_sum integrand0 in\n let dn_0 = trap_sum (Nx.mul_s (Nx.mul integrand0 y2_0) 2.0) in\n let r_nl_t =\n Nx.add_s (Nx.mul_s (Nx.div (Nx.sub_s s2_0 1.0) dn_0) r_nl_f) r_nl_f\n in\n let k_nl = Nx.recip r_nl_t in\n (* Recompute n_eff and C at the tensor R_nl for full differentiability. *)\n let kr = Nx.mul k r_nl_t in\n let y2 = Nx.mul kr kr in\n let gauss = Nx.exp (Nx.neg y2) in\n let integrand = Nx.mul k3pk gauss in\n let s2 = trap_sum integrand in\n let dn = trap_sum (Nx.mul_s (Nx.mul integrand y2) 2.0) in\n let dc =\n trap_sum (Nx.mul (Nx.mul_s integrand 4.0) (Nx.sub y2 (Nx.mul y2 y2)))\n in\n let n_eff = Nx.sub_s dn 3.0 in\n let c_curv = Nx.add (Nx.mul dn dn) (Nx.div dc s2) in\n (k_nl, n_eff, c_curv)\n\nlet nonlinear_power ?(p = default) k z =\n let s = Nx.scalar f64 in\n let pk_lin = linear_power ~p k z in\n let k_nl, n, c = halofit_params p z in\n let n2 = Nx.mul n n in\n let n3 = Nx.mul n2 n in\n let n4 = Nx.mul n3 n in\n let a = Nx.recip (Nx.add_s z 1.0) in\n let om_m = omega_m_a p a in\n let om_de = omega_de_a p a in\n let w = w_of p a in\n let odew1 = Nx.mul om_de (Nx.add_s w 1.0) in\n (* Takahashi et al. 2012 coefficients — all tensor *)\n let a_n =\n Nx.pow (s 10.0)\n (Nx.add\n (Nx.add\n (Nx.add\n (Nx.add\n (Nx.add\n (Nx.add_s (Nx.mul_s n 2.8553) 1.5222)\n (Nx.mul_s n2 2.3706))\n (Nx.mul_s n3 0.9903))\n (Nx.mul_s n4 0.2250))\n (Nx.mul_s c (-0.6038)))\n (Nx.mul_s odew1 0.1749))\n in\n let b_n =\n Nx.pow (s 10.0)\n (Nx.add\n (Nx.add\n (Nx.add\n (Nx.add_s (Nx.mul_s n 0.5864) (-0.5642))\n (Nx.mul_s n2 0.5716))\n (Nx.mul_s c (-1.5474)))\n (Nx.mul_s odew1 0.2279))\n in\n let c_n =\n Nx.pow (s 10.0)\n (Nx.add\n (Nx.add (Nx.add_s (Nx.mul_s n 2.0404) 0.3698) (Nx.mul_s n2 0.8161))\n (Nx.mul_s c 0.5869))\n in\n let gamma_n =\n Nx.add (Nx.add_s (Nx.mul_s n (-0.0843)) 0.1971) (Nx.mul_s c 0.8460)\n in\n let alpha_n =\n Nx.abs\n (Nx.add\n (Nx.add (Nx.add_s (Nx.mul_s n 1.3373) 6.0835) (Nx.mul_s n2 (-0.1959)))\n (Nx.mul_s c (-5.5274)))\n in\n let beta_n =\n Nx.add\n (Nx.add\n (Nx.add\n (Nx.add\n (Nx.add_s (Nx.mul_s n (-0.7354)) 2.0379)\n (Nx.mul_s n2 0.3157))\n (Nx.mul_s n3 1.2490))\n (Nx.mul_s n4 0.3980))\n (Nx.mul_s c (-0.1682))\n in\n let nu_n = Nx.pow (s 10.0) (Nx.add_s (Nx.mul_s n 3.6902) 5.2105) in\n let f1 = Nx.pow om_m (s (-0.0307)) in\n let f2 = Nx.pow om_m (s (-0.0585)) in\n let f3 = Nx.pow om_m (s 0.0743) in\n let y = Nx.div k k_nl in\n (* Δ²_L = k³ P_lin / (2π²) *)\n let d2l =\n Nx.div\n (Nx.mul (Nx.mul k (Nx.mul k k)) pk_lin)\n (s (2.0 *. Float.pi *. Float.pi))\n in\n (* f(y) = y/4 + y²/8 *)\n let fy = Nx.add (Nx.div_s y 4.0) (Nx.div_s (Nx.mul y y) 8.0) in\n (* Quasi-linear term: Δ²_Q *)\n let d2q =\n Nx.mul d2l\n (Nx.mul\n (Nx.div\n (Nx.pow (Nx.add_s d2l 1.0) beta_n)\n (Nx.add_s (Nx.mul d2l alpha_n) 1.0))\n (Nx.exp (Nx.neg fy)))\n in\n (* Halo term: Δ²_H *)\n let three_f1 = Nx.mul_s f1 3.0 in\n let d2h_prime =\n Nx.div\n (Nx.mul a_n (Nx.pow y three_f1))\n (Nx.add_s\n (Nx.add\n (Nx.mul b_n (Nx.pow y f2))\n (Nx.pow (Nx.mul (Nx.mul c_n f3) y) (Nx.sub_s gamma_n 3.0)))\n 1.0)\n in\n let d2h = Nx.div d2h_prime (Nx.add_s (Nx.div nu_n (Nx.mul y y)) 1.0) in\n let d2nl = Nx.add d2q d2h in\n Nx.div (Nx.mul_s d2nl (2.0 *. Float.pi *. Float.pi)) (Nx.mul k (Nx.mul k k))\n\n(* BAO distance measures *)\n\nlet dh ?(p = default) z =\n Unit.Length.of_tensor (Nx.mul_s (Nx.div c_km_s (hubble ~p z)) _mpc_m)\n\nlet dm ?(p = default) z =\n Unit.Length.of_tensor (Nx.mul_s (transverse_comoving_mpc p z) _mpc_m)\n\nlet dv ?(p = default) z =\n let dh_mpc = Nx.div c_km_s (hubble ~p z) in\n let dm_mpc = transverse_comoving_mpc p z in\n let cube = Nx.mul z (Nx.mul dh_mpc (Nx.mul dm_mpc dm_mpc)) in\n Unit.Length.of_tensor (Nx.mul_s (Nx.pow_s cube (1.0 /. 3.0)) _mpc_m)\n\nlet sound_horizon ?(p = default) () =\n let ob = omega_b p in\n let h = Nx.div_s p.h0 100.0 in\n let h2 = Nx.mul h h in\n let w_m = Nx.mul p.omega_m h2 in\n let w_b = Nx.mul ob h2 in\n (* Eisenstein & Hu (1998) Eq. 2–6: sound horizon at drag epoch in Mpc/h *)\n let t27sq = (t_cmb /. 2.7) ** 2.0 in\n let t27_4 = t27sq *. t27sq in\n let z_eq = Nx.div_s (Nx.mul_s w_m 2.50e4) t27_4 in\n let k_eq = Nx.div (Nx.div_s (Nx.mul_s w_m 7.46e-2) t27sq) h in\n let b1_z =\n Nx.mul\n (Nx.pow w_m (Nx.scalar f64 (-0.419)))\n (Nx.add_s (Nx.mul_s (Nx.pow w_m (Nx.scalar f64 0.674)) 0.607) 1.0)\n |> fun x -> Nx.mul_s x 0.313\n in\n let b2_z = Nx.mul_s (Nx.pow w_m (Nx.scalar f64 0.223)) 0.238 in\n let z_d =\n Nx.mul\n (Nx.div\n (Nx.mul_s (Nx.pow w_m (Nx.scalar f64 0.251)) 1291.0)\n (Nx.add_s (Nx.mul_s (Nx.pow w_m (Nx.scalar f64 0.828)) 0.659) 1.0))\n (Nx.add_s (Nx.mul b1_z (Nx.pow w_b b2_z)) 1.0)\n in\n let r_d =\n Nx.mul (Nx.div_s (Nx.mul_s w_b 31.5) t27_4) (Nx.div (Nx.scalar f64 1e3) z_d)\n in\n let r_eq =\n Nx.mul\n (Nx.div_s (Nx.mul_s w_b 31.5) t27_4)\n (Nx.div (Nx.scalar f64 1e3) z_eq)\n in\n (* Eq. 6 from Eisenstein & Hu: sound horizon in Mpc/h *)\n let sh_d =\n Nx.mul\n (Nx.mul\n (Nx.div (Nx.scalar f64 2.0) (Nx.mul_s k_eq 3.0))\n (Nx.sqrt (Nx.div (Nx.scalar f64 6.0) r_eq)))\n (Nx.log\n (Nx.div\n (Nx.add (Nx.sqrt (Nx.add_s r_d 1.0)) (Nx.sqrt (Nx.add r_eq r_d)))\n (Nx.add_s (Nx.sqrt r_eq) 1.0)))\n in\n (* sh_d is in Mpc/h, convert to Mpc then to metres *)\n let rs_mpc = Nx.div sh_d h in\n Unit.Length.of_tensor (Nx.mul_s rs_mpc _mpc_m)\n"
},
{
"path": "dev/umbra/lib/cosmo.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Cosmology for {e Λ}CDM, wCDM, and w0waCDM universes.\n\n Computes distances, growth factors, and matter power spectra. Supports flat\n and non-flat {e Λ}CDM, wCDM, and w0waCDM cosmologies through a single\n parameter type. All functions are differentiable through Rune.\n\n {[\n let z = Nx.scalar Nx.float64 0.5 in\n let dl = Cosmo.luminosity_distance z in\n let dl_mpc = Unit.Length.in_mpc dl\n ]}\n\n Power spectrum functions require [omega_b], [n_s], and [sigma8] to be set\n via {!set} or by using a preset like {!planck18}. *)\n\n(** {1:params Parameters} *)\n\ntype params\n(** The type for cosmological parameters. Subsumes flat {e Λ}CDM, non-flat\n {e Λ}CDM, wCDM, and w0waCDM. *)\n\n(** {2:float_constructors Float constructors}\n\n Create parameters from plain floats. *)\n\nval flat_lcdm : h0:float -> omega_m:float -> params\n(** [flat_lcdm ~h0 ~omega_m] is flat {e Λ}CDM with {e Ω}{_ L}[ = 1 - omega_m].\n\n Raises [Invalid_argument] if [h0 <= 0] or [omega_m < 0]. *)\n\nval lcdm : h0:float -> omega_m:float -> omega_l:float -> params\n(** [lcdm ~h0 ~omega_m ~omega_l] is {e Λ}CDM with curvature\n {e Ω}{_ k}[ = 1 - omega_m - omega_l].\n\n Raises [Invalid_argument] if [h0 <= 0]. *)\n\nval wcdm :\n h0:float -> omega_m:float -> ?omega_l:float -> w0:float -> unit -> params\n(** [wcdm ~h0 ~omega_m ~w0 ()] is wCDM with constant dark energy equation of\n state [w0]. [omega_l] defaults to [1 - omega_m] (flat). *)\n\nval w0wacdm :\n h0:float ->\n omega_m:float ->\n ?omega_l:float ->\n w0:float ->\n wa:float ->\n unit ->\n params\n(** [w0wacdm ~h0 ~omega_m ~w0 ~wa ()] is the CPL parameterization\n [w(z) = w0 + wa * z/(1+z)]. [omega_l] defaults to [1 - omega_m] (flat). *)\n\n(** {2:tensor_constructors Tensor constructors}\n\n Create parameters from Nx scalar tensors for differentiable construction. *)\n\nval create_flat_lcdm : h0:Nx.float64_t -> omega_m:Nx.float64_t -> params\n\nval create_lcdm :\n h0:Nx.float64_t -> omega_m:Nx.float64_t -> omega_l:Nx.float64_t -> params\n\nval create_wcdm :\n h0:Nx.float64_t ->\n omega_m:Nx.float64_t ->\n ?omega_l:Nx.float64_t ->\n w0:Nx.float64_t ->\n unit ->\n params\n\nval create_w0wacdm :\n h0:Nx.float64_t ->\n omega_m:Nx.float64_t ->\n ?omega_l:Nx.float64_t ->\n w0:Nx.float64_t ->\n wa:Nx.float64_t ->\n unit ->\n params\n\n(** {2:accessors Accessors} *)\n\nval h0 : params -> Nx.float64_t\n(** [h0 p] is the Hubble constant H{_ 0} in km s{^ -1} Mpc{^ -1}. *)\n\nval omega_m : params -> Nx.float64_t\n(** [omega_m p] is the matter density parameter {e Ω}{_ m}. *)\n\nval omega_l : params -> Nx.float64_t\n(** [omega_l p] is the dark energy density parameter {e Ω}{_ Λ}. *)\n\nval omega_k : params -> Nx.float64_t\n(** [omega_k p] is the curvature density parameter {e Ω}{_ k}[ = 1 - Ω_m - Ω_Λ].\n*)\n\nval w0 : params -> Nx.float64_t\n(** [w0 p] is the dark energy equation of state parameter w{_ 0}. *)\n\nval wa : params -> Nx.float64_t\n(** [wa p] is the CPL time-varying dark energy parameter w{_ a}. *)\n\nval omega_b : params -> Nx.float64_t\n(** [omega_b p] is the baryon density parameter {e Ω}{_ b}.\n\n Raises [Invalid_argument] if not set. *)\n\nval n_s : params -> Nx.float64_t\n(** [n_s p] is the primordial spectral index n{_ s}.\n\n Raises [Invalid_argument] if not set. *)\n\nval sigma8 : params -> Nx.float64_t\n(** [sigma8 p] is the amplitude of matter fluctuations {e σ}{_ 8}.\n\n Raises [Invalid_argument] if not set. *)\n\n(** {2:set Setting power spectrum parameters} *)\n\nval set : ?omega_b:float -> ?n_s:float -> ?sigma8:float -> params -> params\n(** [set ~omega_b ~n_s ~sigma8 p] is [p] with the given power spectrum\n parameters set. Unspecified parameters retain their previous value. *)\n\nval set_t :\n ?h0:Nx.float64_t ->\n ?omega_m:Nx.float64_t ->\n ?omega_l:Nx.float64_t ->\n ?omega_b:Nx.float64_t ->\n ?n_s:Nx.float64_t ->\n ?sigma8:Nx.float64_t ->\n params ->\n params\n(** [set_t] is like {!set} but takes Nx scalar tensors for differentiable\n construction. Recomputes {e Ω}{_ k} when [omega_m] or [omega_l] changes. *)\n\n(** {2:presets Presets} *)\n\nval default : params\n(** [default] is flat {e Λ}CDM with [h0 = 70], [omega_m = 0.3]. *)\n\nval planck18 : params\n(** [planck18] is Planck 2018 flat {e Λ}CDM: [h0 = 67.66], [omega_m = 0.3111],\n [omega_b = 0.0490], [n_s = 0.9665], [sigma8 = 0.8102]. *)\n\nval planck15 : params\n(** [planck15] is Planck 2015 flat {e Λ}CDM: [h0 = 67.74], [omega_m = 0.3075],\n [omega_b = 0.0486], [n_s = 0.9667], [sigma8 = 0.8159]. *)\n\nval wmap9 : params\n(** [wmap9] is WMAP9 flat {e Λ}CDM: [h0 = 69.32], [omega_m = 0.2865],\n [omega_b = 0.0463], [n_s = 0.9608], [sigma8 = 0.820]. *)\n\n(** {1:e_z Hubble parameter} *)\n\nval e_of : params -> Nx.float64_t -> Nx.float64_t\n(** [e_of p z] is E(z) = H(z)/H{_ 0} at redshift [z]. Fully differentiable\n through Rune. *)\n\nval hubble : ?p:params -> Nx.float64_t -> Nx.float64_t\n(** [hubble z] is H(z) in km s{^ -1} Mpc{^ -1}. [p] defaults to {!default}. *)\n\nval critical_density : ?p:params -> Nx.float64_t -> Nx.float64_t\n(** [critical_density z] is the critical density {e rho}{_ c}(z) in kg m{^ -3}.\n [p] defaults to {!default}. *)\n\n(** {1:distances Distances} *)\n\nval comoving_distance : ?p:params -> Nx.float64_t -> Unit.length Unit.t\n(** [comoving_distance z] is the line-of-sight comoving distance at redshift\n [z]. [p] defaults to {!default}. *)\n\nval luminosity_distance : ?p:params -> Nx.float64_t -> Unit.length Unit.t\n(** [luminosity_distance z] is the luminosity distance at redshift [z]. For\n non-flat models, applies the curvature correction via the transverse\n comoving distance. [p] defaults to {!default}. *)\n\nval angular_diameter_distance : ?p:params -> Nx.float64_t -> Unit.length Unit.t\n(** [angular_diameter_distance z] is the angular diameter distance at redshift\n [z]. [p] defaults to {!default}. *)\n\nval distance_modulus : ?p:params -> Nx.float64_t -> Nx.float64_t\n(** [distance_modulus z] is the distance modulus\n {e mu}[ = 5 log10(d_L / Mpc) + 25]. [p] defaults to {!default}. *)\n\n(** {1:angular Angular scale} *)\n\nval angular_size :\n ?p:params -> z:Nx.float64_t -> Unit.length Unit.t -> Unit.angle Unit.t\n(** [angular_size ~z length] is the angular size of [length] at redshift [z]\n under the small-angle approximation [{e theta} = l / d_A]. [p] defaults to\n {!default}. *)\n\nval physical_size :\n ?p:params -> z:Nx.float64_t -> Unit.angle Unit.t -> Unit.length Unit.t\n(** [physical_size ~z angle] is the physical size subtended by [angle] at\n redshift [z] under the small-angle approximation [l = {e theta} * d_A]. [p]\n defaults to {!default}. *)\n\n(** {1:times Cosmic times} *)\n\nval lookback_time : ?p:params -> Nx.float64_t -> Unit.time Unit.t\n(** [lookback_time z] is the lookback time to redshift [z]. [p] defaults to\n {!default}. *)\n\nval age : ?p:params -> Nx.float64_t -> Unit.time Unit.t\n(** [age z] is the age of the universe at redshift [z].\n\n Integrates from [z] to [z = 1000]. This approximation is accurate to ~0.1%\n for late-time cosmology ([z < 10]) but omits the radiation era and is not\n suitable for CMB-epoch calculations. [p] defaults to {!default}. *)\n\n(** {1:inverse Inverse lookup} *)\n\nval z_at_value :\n ?p:params ->\n ?zmin:float ->\n ?zmax:float ->\n ?xtol:float ->\n (p:params -> Nx.float64_t -> Nx.float64_t) ->\n Nx.float64_t ->\n Nx.float64_t\n(** [z_at_value f target] finds the redshift [z] where [f ~p z = target] using\n Brent's method. [f] must be a monotonic function of redshift.\n\n For distance functions, unwrap the unit first:\n {[\n z_at_value\n (fun ~p z -> Unit.Length.in_mpc (Cosmo.comoving_distance ~p z))\n target\n ]}\n\n [zmin] defaults to [1e-8]. [zmax] defaults to [1000.0]. [xtol] defaults to\n [1e-8].\n\n {b Warning.} Not differentiable (iterative root-finding).\n\n Raises [Invalid_argument] if [target] is outside [[f(zmin), f(zmax)]]. *)\n\n(** {1:bao BAO distance measures} *)\n\nval dh : ?p:params -> Nx.float64_t -> Unit.length Unit.t\n(** [dh z] is the Hubble distance D{_ H}(z) = c / H(z). [p] defaults to\n {!default}. *)\n\nval dm : ?p:params -> Nx.float64_t -> Unit.length Unit.t\n(** [dm z] is the comoving transverse distance D{_ M}(z). Equal to\n {!comoving_distance} for flat cosmologies; includes curvature correction\n otherwise. [p] defaults to {!default}. *)\n\nval dv : ?p:params -> Nx.float64_t -> Unit.length Unit.t\n(** [dv z] is the volume-averaged BAO distance D{_ V}(z) = (z D{_ H}(z)\n D{_ M}{^ 2}(z)){^ 1/3}. [p] defaults to {!default}. *)\n\nval sound_horizon : ?p:params -> unit -> Unit.length Unit.t\n(** [sound_horizon ()] is the comoving sound horizon at the drag epoch\n r{_ s}(z{_ drag}), using the Eisenstein & Hu (1998) fitting formulae for\n z{_ drag} and the sound horizon integral.\n\n Raises [Invalid_argument] if [omega_b] is not set in [p]. [p] defaults to\n {!default}. *)\n\n(** {1:growth Structure growth} *)\n\nval growth_factor : ?p:params -> Nx.float64_t -> Nx.float64_t\n(** [growth_factor z] is the linear growth factor D(z), normalized to D(0) = 1.\n Computed via the integral form D(a) {e ∝} E(a) {e ∫}{_ 0}{^ a} da' /\n (a'{^ 3} E{^ 3}(a')).\n\n Does not require [omega_b], [n_s], or [sigma8]. [p] defaults to {!default}.\n*)\n\nval growth_rate : ?p:params -> Nx.float64_t -> Nx.float64_t\n(** [growth_rate z] is the linear growth rate f(z) = d ln D / d ln a, computed\n from the exact derivative of the integral-form growth factor.\n\n [p] defaults to {!default}. *)\n\n(** {1:power Matter power spectrum}\n\n All power spectrum functions require [omega_b], [n_s], and [sigma8] to be\n set in the parameters. Use {!set} or a preset like {!planck18}.\n\n Wavenumbers [k] are in h/Mpc. Power spectra are in (Mpc/h){^ 3}. *)\n\nval linear_power : ?p:params -> Nx.float64_t -> Nx.float64_t -> Nx.float64_t\n(** [linear_power ~p k z] is the linear matter power spectrum P(k, z). Uses the\n Eisenstein & Hu (1998) transfer function with baryon oscillations and\n {e σ}{_ 8} normalization.\n\n Raises [Invalid_argument] if [omega_b], [n_s], or [sigma8] are not set. *)\n\nval nonlinear_power : ?p:params -> Nx.float64_t -> Nx.float64_t -> Nx.float64_t\n(** [nonlinear_power ~p k z] is the nonlinear matter power spectrum via the\n Halofit fitting formula (Takahashi et al. 2012).\n\n {b Warning.} The nonlinear scale k{_ nl} is found by float-level\n root-finding; gradients do not flow through it. The mapping from k{_ nl} to\n P{_ nl}(k) is differentiable.\n\n Raises [Invalid_argument] if [omega_b], [n_s], or [sigma8] are not set. *)\n"
},
{
"path": "dev/umbra/lib/dune",
"content": "(library\n (name umbra)\n (public_name umbra)\n (private_modules kdtree filter_data vega_data)\n (libraries nx unix))\n"
},
{
"path": "dev/umbra/lib/extinction.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet f64 = Nx.float64\n\n(* Extinction law: wavelength in metres → A_λ/A_V *)\ntype law = wavelength:Nx.float64_t -> Nx.float64_t\n\n(* Horner evaluation: c0 + y*(c1 + y*(c2 + ...)) *)\nlet horner y coeffs =\n let n = Array.length coeffs in\n let acc = ref (Nx.scalar f64 coeffs.(n - 1)) in\n for i = n - 2 downto 0 do\n acc := Nx.add_s (Nx.mul y !acc) coeffs.(i)\n done;\n !acc\n\n(* Shared CCM89/O'Donnell94 implementation parameterized by R_V. Only the\n optical a/b polynomial coefficients differ between the two laws; IR and UV\n regions are identical. Uses Nx.where for differentiable piecewise selection.\n Valid for 0.125–3.5 μm (x = 0.3–8.0 μm⁻¹). *)\nlet ccm89_impl a_opt b_opt ~rv ~wavelength =\n (* Convert wavelength (m) to inverse microns *)\n let x = Nx.div (Nx.scalar f64 1e-6) wavelength in\n (* Infrared: 0.3 ≤ x < 1.1 *)\n let a_ir = Nx.mul_s (Nx.pow_s x 1.61) 0.574 in\n let b_ir = Nx.mul_s (Nx.pow_s x 1.61) (-0.527) in\n (* Optical/NIR: 1.1 ≤ x ≤ 3.3, polynomial in (x - 1.82) *)\n let y = Nx.sub_s x 1.82 in\n let a_o = horner y a_opt in\n let b_o = horner y b_opt in\n (* UV: 3.3 < x ≤ 8.0 *)\n let fa =\n Nx.where (Nx.greater_equal_s x 5.9)\n (Nx.add\n (Nx.mul_s (Nx.square (Nx.sub_s x 5.9)) (-0.04473))\n (Nx.mul_s (Nx.pow_s (Nx.sub_s x 5.9) 3.0) (-0.009779)))\n (Nx.scalar f64 0.0)\n in\n let fb =\n Nx.where (Nx.greater_equal_s x 5.9)\n (Nx.add\n (Nx.mul_s (Nx.square (Nx.sub_s x 5.9)) 0.2130)\n (Nx.mul_s (Nx.pow_s (Nx.sub_s x 5.9) 3.0) 0.1207))\n (Nx.scalar f64 0.0)\n in\n (* a(x) = 1.752 - 0.316*x - 0.104/((x-4.67)² + 0.341) + F_a *)\n let a_uv_base = Nx.add_s (Nx.mul_s x (-0.316)) 1.752 in\n let bump_a =\n Nx.div (Nx.scalar f64 (-0.104))\n (Nx.add (Nx.square (Nx.sub_s x 4.67)) (Nx.scalar f64 0.341))\n in\n let a_uv = Nx.add (Nx.add a_uv_base bump_a) fa in\n (* b(x) = -3.090 + 1.825*x + 1.206/((x-4.62)² + 0.263) + F_b *)\n let b_uv_base = Nx.add_s (Nx.mul_s x 1.825) (-3.090) in\n let bump_b =\n Nx.div (Nx.scalar f64 1.206)\n (Nx.add (Nx.square (Nx.sub_s x 4.62)) (Nx.scalar f64 0.263))\n in\n let b_uv = Nx.add (Nx.add b_uv_base bump_b) fb in\n (* Piecewise selection using Nx.where *)\n let ir_mask = Nx.less_s x 1.1 in\n let uv_mask = Nx.greater_s x 3.3 in\n let a = Nx.where ir_mask a_ir (Nx.where uv_mask a_uv a_o) in\n let b = Nx.where ir_mask b_ir (Nx.where uv_mask b_uv b_o) in\n (* A_λ/A_V = a(x) + b(x)/R_V *)\n Nx.add a (Nx.div b rv)\n\n(* CCM89: Cardelli, Clayton & Mathis 1989, ApJ 345, 245 — optical\n coefficients *)\nlet ccm89_a =\n [| 1.0; 0.17699; -0.50447; -0.02427; 0.72085; 0.01979; -0.77530; 0.32999 |]\n\nlet ccm89_b =\n [| 0.0; 1.41338; 2.28305; 1.07233; -5.38434; -0.62251; 5.30260; -2.09002 |]\n\nlet ccm89 ~rv = fun ~wavelength -> ccm89_impl ccm89_a ccm89_b ~rv ~wavelength\n\n(* O'Donnell 1994, ApJ 422, 158 — revised optical coefficients *)\nlet od94_a = [| 1.0; 0.104; -0.609; 0.701; -1.221; 0.700; -0.048; -0.091 |]\nlet od94_b = [| 0.0; 1.952; 2.908; -3.989; 7.985; -5.002; -0.478; 1.149 |]\nlet odonnell94 ~rv = fun ~wavelength -> ccm89_impl od94_a od94_b ~rv ~wavelength\n\n(* Calzetti 2000: Calzetti et al. 2000, ApJ 533, 682. Starburst attenuation law.\n Fixed R_V = 4.05. Valid 0.12–2.2 μm. *)\nlet calzetti00 =\n fun ~wavelength ->\n let lam_um = Nx.mul_s wavelength 1e6 in\n let rv = 4.05 in\n (* Blue: 0.12–0.63 μm k'(λ) = 2.659 * (-2.156 + 1.509/λ - 0.198/λ² + 0.011/λ³)\n + R_V *)\n let k_blue =\n Nx.add_s\n (Nx.mul_s\n (Nx.add_s\n (Nx.add\n (Nx.mul_s (Nx.recip lam_um) 1.509)\n (Nx.add\n (Nx.mul_s (Nx.pow_s lam_um (-2.0)) (-0.198))\n (Nx.mul_s (Nx.pow_s lam_um (-3.0)) 0.011)))\n (-2.156))\n 2.659)\n rv\n in\n (* Red: 0.63–2.2 μm k'(λ) = 2.659 * (-1.857 + 1.040/λ) + R_V *)\n let k_red =\n Nx.add_s\n (Nx.mul_s (Nx.add_s (Nx.mul_s (Nx.recip lam_um) 1.040) (-1.857)) 2.659)\n rv\n in\n let blue_mask = Nx.less_s lam_um 0.63 in\n let k = Nx.where blue_mask k_blue k_red in\n (* A_λ/A_V = k'(λ) / R_V *)\n Nx.div_s k rv\n\n(* Fitzpatrick 1999: Fitzpatrick 1999, PASP 111, 63. R_V-dependent extinction\n using cubic spline for optical/NIR and Fitzpatrick & Massa parameterization\n for UV. Valid 0.1–3.5 μm. *)\n\n(* FM UV parameters (fixed) *)\nlet f99_x0_sq = 4.596 *. 4.596\nlet f99_gamma_sq = 0.99 *. 0.99\nlet f99_c3 = 3.23\nlet f99_c4 = 0.41\nlet f99_c5 = 5.9\n\n(* Spline anchor x-values (inverse microns) *)\nlet f99_xk =\n [|\n 0.;\n 1e4 /. 26500.;\n 1e4 /. 12200.;\n 1e4 /. 6000.;\n 1e4 /. 5470.;\n 1e4 /. 4670.;\n 1e4 /. 4110.;\n 1e4 /. 2700.;\n 1e4 /. 2600.;\n |]\n\nlet f99_hk = Array.init 8 (fun i -> f99_xk.(i + 1) -. f99_xk.(i))\n\n(* Drude profile at a fixed x-value *)\nlet f99_drude x =\n let x2 = x *. x in\n let y = x2 -. f99_x0_sq in\n x2 /. ((y *. y) +. (x2 *. f99_gamma_sq))\n\n(* Precompute spline basis matrix M (7×9): maps 9 anchor y-values to 7 interior\n second derivatives. Natural boundary conditions: m[0] = m[8] = 0.\n\n The tridiagonal system Am = Dy is solved offline; M = A⁻¹D is stored. At\n runtime m[j] = Σ M[j][i] y[i] — a weighted sum of Nx scalars, fully\n differentiable through Rune. *)\nlet f99_basis =\n let n = 7 in\n let h = f99_hk in\n (* Right-hand side matrix D (7×9) *)\n let d_mat =\n Array.init n (fun j ->\n Array.init 9 (fun i ->\n if i = j then 6.0 /. h.(j)\n else if i = j + 1 then ~-.((6.0 /. h.(j + 1)) +. (6.0 /. h.(j)))\n else if i = j + 2 then 6.0 /. h.(j + 1)\n else 0.0))\n in\n (* Tridiagonal A: diag, sub, sup *)\n let diag = Array.init n (fun j -> 2.0 *. (h.(j) +. h.(j + 1))) in\n let sub j = h.(j) in\n let sup j = h.(j + 1) in\n (* Solve A X_col = D_col for each of 9 columns via Thomas algorithm *)\n let m = Array.init n (fun _ -> Array.make 9 0.0) in\n for col = 0 to 8 do\n let b = Array.init n (fun j -> d_mat.(j).(col)) in\n let c = Array.make n 0.0 in\n let d = Array.make n 0.0 in\n c.(0) <- sup 0 /. diag.(0);\n d.(0) <- b.(0) /. diag.(0);\n for i = 1 to n - 1 do\n let w = diag.(i) -. (sub i *. c.(i - 1)) in\n c.(i) <- (if i < n - 1 then sup i /. w else 0.0);\n d.(i) <- (b.(i) -. (sub i *. d.(i - 1))) /. w\n done;\n m.(n - 1).(col) <- d.(n - 1);\n for i = n - 2 downto 0 do\n m.(i).(col) <- d.(i) -. (c.(i) *. m.(i + 1).(col))\n done\n done;\n m\n\n(* Evaluate a cubic spline piece on [xk, xk1] at tensor x. mk and mk1 are second\n derivatives (Nx scalars); yk, yk1 are y-values. *)\nlet f99_eval_piece hk xk yk yk1 mk mk1 x =\n let a = yk in\n let c = Nx.mul_s mk 0.5 in\n let d = Nx.div_s (Nx.sub mk1 mk) (6.0 *. hk) in\n let b =\n Nx.sub\n (Nx.div_s (Nx.sub yk1 yk) hk)\n (Nx.mul_s (Nx.add (Nx.mul_s mk 2.0) mk1) (hk /. 6.0))\n in\n let t = Nx.sub_s x xk in\n Nx.add a (Nx.mul t (Nx.add b (Nx.mul t (Nx.add c (Nx.mul t d)))))\n\nlet fitzpatrick99 ~rv =\n let rv2 = Nx.mul rv rv in\n let rv3 = Nx.mul rv2 rv in\n let rv4 = Nx.mul rv2 rv2 in\n (* FM UV c1, c2 — computed once, used for anchor y-values and the closure *)\n let c2_uv = Nx.add_s (Nx.mul_s (Nx.recip rv) 4.717) (-0.824) in\n let c1_uv = Nx.sub (Nx.scalar f64 2.030) (Nx.mul_s c2_uv 3.007) in\n let uv_anchor xk =\n Nx.add c1_uv (Nx.add_s (Nx.mul_s c2_uv xk) (f99_c3 *. f99_drude xk))\n in\n (* 9 anchor E(λ-V)/E(B-V) values *)\n let y =\n [|\n Nx.neg rv;\n Nx.sub (Nx.mul_s rv (0.26469 /. 3.1)) rv;\n Nx.sub (Nx.mul_s rv (0.82925 /. 3.1)) rv;\n Nx.sub\n (Nx.add\n (Nx.add_s (Nx.mul_s rv 1.00270) (-0.422809))\n (Nx.mul_s rv2 2.13572e-04))\n rv;\n Nx.sub\n (Nx.add\n (Nx.add_s (Nx.mul_s rv 1.00216) (-5.13540e-02))\n (Nx.mul_s rv2 (-7.35778e-05)))\n rv;\n Nx.sub\n (Nx.add\n (Nx.add_s (Nx.mul_s rv 1.00184) 0.700127)\n (Nx.mul_s rv2 (-3.32598e-05)))\n rv;\n Nx.sub\n (Nx.add\n (Nx.add\n (Nx.add\n (Nx.add_s (Nx.mul_s rv 1.01707) 1.19456)\n (Nx.mul_s rv2 (-5.46959e-03)))\n (Nx.mul_s rv3 7.97809e-04))\n (Nx.mul_s rv4 (-4.45636e-05)))\n rv;\n (* UV anchors from FM parameterization *)\n uv_anchor f99_xk.(7);\n uv_anchor f99_xk.(8);\n |]\n in\n (* Second derivatives m[0..8]: m[0] = m[8] = 0, m[1..7] from basis matrix *)\n let zero = Nx.scalar f64 0.0 in\n let m2 = Array.make 9 zero in\n for j = 0 to 6 do\n let acc = ref zero in\n for i = 0 to 8 do\n acc := Nx.add !acc (Nx.mul_s y.(i) f99_basis.(j).(i))\n done;\n m2.(j + 1) <- !acc\n done;\n (* Precompute spline piece coefficients for intervals 0..6 *)\n let h = f99_hk in\n let pieces =\n Array.init 7 (fun k ->\n let hk = h.(k) in\n let yk = y.(k) in\n let yk1 = y.(k + 1) in\n let mk = m2.(k) in\n let mk1 = m2.(k + 1) in\n (hk, f99_xk.(k), yk, yk1, mk, mk1))\n in\n fun ~wavelength ->\n (* Convert wavelength (m) to inverse microns *)\n let x = Nx.div (Nx.scalar f64 1e-6) wavelength in\n (* Evaluate spline for each interval *)\n let eval k =\n let hk, xk, yk, yk1, mk, mk1 = pieces.(k) in\n f99_eval_piece hk xk yk yk1 mk mk1 x\n in\n let s0 = eval 0 in\n let s1 = eval 1 in\n let s2 = eval 2 in\n let s3 = eval 3 in\n let s4 = eval 4 in\n let s5 = eval 5 in\n let s6 = eval 6 in\n let opt_nir =\n Nx.where\n (Nx.less_s x f99_xk.(1))\n s0\n (Nx.where\n (Nx.less_s x f99_xk.(2))\n s1\n (Nx.where\n (Nx.less_s x f99_xk.(3))\n s2\n (Nx.where\n (Nx.less_s x f99_xk.(4))\n s3\n (Nx.where\n (Nx.less_s x f99_xk.(5))\n s4\n (Nx.where (Nx.less_s x f99_xk.(6)) s5 s6)))))\n in\n (* UV: FM parameterization for x ≥ 1e4/2700 *)\n let x2 = Nx.square x in\n let y_bump = Nx.sub x2 (Nx.scalar f64 f99_x0_sq) in\n let drude =\n Nx.div x2 (Nx.add (Nx.mul y_bump y_bump) (Nx.mul_s x2 f99_gamma_sq))\n in\n let fuv =\n Nx.where\n (Nx.greater_equal_s x f99_c5)\n (let dx = Nx.sub_s x f99_c5 in\n let dx2 = Nx.square dx in\n Nx.add (Nx.mul_s dx2 0.5392) (Nx.mul_s (Nx.mul dx2 dx) 0.05644))\n (Nx.scalar f64 0.0)\n in\n let k_uv =\n Nx.add c1_uv\n (Nx.add (Nx.mul c2_uv x)\n (Nx.add (Nx.mul_s drude f99_c3) (Nx.mul_s fuv f99_c4)))\n in\n (* Select optical/NIR vs UV *)\n let e_over_ebv = Nx.where (Nx.less_s x f99_xk.(7)) opt_nir k_uv in\n (* A(λ)/A(V) = E(λ-V)/E(B-V) / R_V + 1 *)\n Nx.add_s (Nx.div e_over_ebv rv) 1.0\n\nlet curve law ~wavelength = law ~wavelength:(Unit.Length.to_tensor wavelength)\nlet ln10_over_2_5 = Float.log 10.0 *. 0.4\n\nlet scale_flux sign law ~av spectrum =\n let wave_m = Unit.Length.to_tensor (Spectrum.wavelength spectrum) in\n let a_lambda = Nx.mul (law ~wavelength:wave_m) av in\n let factor = Nx.exp (Nx.mul_s a_lambda (sign *. ln10_over_2_5)) in\n Spectrum.scale factor spectrum\n\nlet apply law ~av spectrum = scale_flux (-1.0) law ~av spectrum\nlet unredden law ~av spectrum = scale_flux 1.0 law ~av spectrum\n"
},
{
"path": "dev/umbra/lib/extinction.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Dust extinction laws.\n\n Extinction laws describe how interstellar dust attenuates and reddens light\n as a function of wavelength. A {!law} maps wavelength to the normalised\n extinction curve A{_ lambda} / A{_ V}.\n\n {!apply} and {!unredden} are differentiable through Rune with respect to\n [av]. The extinction curve evaluation itself (law constructors and {!curve})\n is not differentiable (scalar-level polynomial and spline evaluation). *)\n\n(** {1:types Types} *)\n\ntype law\n(** The type for extinction laws. *)\n\n(** {1:laws Standard laws} *)\n\nval ccm89 : rv:Nx.float64_t -> law\n(** [ccm89 ~rv] is the\n {{:https://ui.adsabs.harvard.edu/abs/1989ApJ...345..245C}Cardelli, Clayton &\n Mathis (1989)} Milky Way extinction law. [rv] is the total-to-selective\n extinction ratio R{_ V} (typically 3.1).\n\n Valid for 0.125--3.5 {e mu}m (0.3--8.0 {e mu}m{^ -1}). Values outside this\n range are extrapolations. *)\n\nval fitzpatrick99 : rv:Nx.float64_t -> law\n(** [fitzpatrick99 ~rv] is the\n {{:https://ui.adsabs.harvard.edu/abs/1999PASP..111...63F}Fitzpatrick (1999)}\n R{_ V}-dependent Milky Way extinction law. Uses a cubic spline for\n optical/NIR and the Fitzpatrick & Massa UV parameterization.\n\n Valid for 0.1--3.5 {e mu}m (0.3--10.0 {e mu}m{^ -1}). *)\n\nval odonnell94 : rv:Nx.float64_t -> law\n(** [odonnell94 ~rv] is the\n {{:https://ui.adsabs.harvard.edu/abs/1994ApJ...422..158O}O'Donnell (1994)}\n Milky Way extinction law. Identical to {!ccm89} except for revised optical\n coefficients (1.1--3.3 {e mu}m{^ -1}).\n\n Valid for 0.125--3.5 {e mu}m. *)\n\nval calzetti00 : law\n(** [calzetti00] is the\n {{:https://ui.adsabs.harvard.edu/abs/2000ApJ...533..682C}Calzetti et al.\n (2000)} starburst attenuation law with fixed R{_ V} = 4.05.\n\n Valid for 0.12--2.2 {e mu}m. Values outside this range are extrapolations.\n*)\n\n(** {1:evaluation Evaluation} *)\n\nval curve : law -> wavelength:Unit.length Unit.t -> Nx.float64_t\n(** [curve law ~wavelength] is A{_ lambda} / A{_ V} at the given wavelengths.\n Not differentiable. *)\n\n(** {1:application Application} *)\n\nval apply : law -> av:Nx.float64_t -> 'a Spectrum.t -> 'a Spectrum.t\n(** [apply law ~av spectrum] reddens [spectrum] by applying [av] magnitudes of\n V-band extinction. The spectral kind is preserved. Differentiable through\n Rune with respect to [av] and the spectrum values. *)\n\nval unredden : law -> av:Nx.float64_t -> 'a Spectrum.t -> 'a Spectrum.t\n(** [unredden law ~av spectrum] de-reddens [spectrum] by removing [av]\n magnitudes of V-band extinction. The spectral kind is preserved.\n Differentiable through Rune with respect to [av] and the spectrum values. *)\n"
},
{
"path": "dev/umbra/lib/filter_data.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n[@@@ocamlformat \"disable\"]\n\n(* Filter transmission curves from the SVO Filter Profile Service.\n http://svo2.cab.inta-csic.es/theory/fps/\n Wavelengths in Angstroms, throughput dimensionless. *)\n\n(* SDSS *)\n\nlet sdss_u_wave =\n [|\n 2980.0; 3005.0; 3030.0; 3055.0; 3080.0; 3105.0;\n 3130.0; 3155.0; 3180.0; 3205.0; 3230.0; 3255.0;\n 3280.0; 3305.0; 3330.0; 3355.0; 3380.0; 3405.0;\n 3430.0; 3455.0; 3480.0; 3505.0; 3530.0; 3555.0;\n 3580.0; 3605.0; 3630.0; 3655.0; 3680.0; 3705.0;\n 3730.0; 3755.0; 3780.0; 3805.0; 3830.0; 3855.0;\n 3880.0; 3905.0; 3930.0; 3955.0; 3980.0; 4005.0;\n 4030.0; 4055.0; 4080.0; 4105.0; 4130.0\n |]\n\nlet sdss_u_thru =\n [|\n 0.0; 0.0001; 0.0005; 0.0013; 0.0026; 0.0052;\n 0.0093; 0.0161; 0.024; 0.0323; 0.0405; 0.0485;\n 0.0561; 0.0634; 0.07; 0.0756; 0.0803; 0.0848;\n 0.0883; 0.0917; 0.0959; 0.1001; 0.1029; 0.1044;\n 0.1053; 0.1063; 0.1075; 0.1085; 0.1084; 0.1064;\n 0.1024; 0.0966; 0.0887; 0.0787; 0.0672; 0.0549;\n 0.0413; 0.0268; 0.0145; 0.0075; 0.0042; 0.0022;\n 0.001; 0.0006; 0.0004; 0.0002; 0.0\n |]\n\n(* SDSS *)\n\nlet sdss_g_wave =\n [|\n 3630.0; 3655.0; 3680.0; 3705.0; 3730.0; 3755.0;\n 3780.0; 3805.0; 3830.0; 3855.0; 3880.0; 3905.0;\n 3930.0; 3955.0; 3980.0; 4005.0; 4030.0; 4055.0;\n 4080.0; 4105.0; 4130.0; 4155.0; 4180.0; 4205.0;\n 4230.0; 4255.0; 4280.0; 4305.0; 4330.0; 4355.0;\n 4380.0; 4405.0; 4430.0; 4455.0; 4480.0; 4505.0;\n 4530.0; 4555.0; 4580.0; 4605.0; 4630.0; 4655.0;\n 4680.0; 4705.0; 4730.0; 4755.0; 4780.0; 4805.0;\n 4830.0; 4855.0; 4880.0; 4905.0; 4930.0; 4955.0;\n 4980.0; 5005.0; 5030.0; 5055.0; 5080.0; 5105.0;\n 5130.0; 5155.0; 5180.0; 5205.0; 5230.0; 5255.0;\n 5280.0; 5305.0; 5330.0; 5355.0; 5380.0; 5405.0;\n 5430.0; 5455.0; 5480.0; 5505.0; 5530.0; 5555.0;\n 5580.0; 5605.0; 5630.0; 5655.0; 5680.0; 5705.0;\n 5730.0; 5755.0; 5780.0; 5805.0; 5830.0\n |]\n\nlet sdss_g_thru =\n [|\n 0.0; 0.0003; 0.0008; 0.0013; 0.0019; 0.0024;\n 0.0034; 0.0055; 0.0103; 0.0194; 0.0326; 0.0492;\n 0.0686; 0.09; 0.1123; 0.1342; 0.1545; 0.1722;\n 0.1873; 0.2003; 0.2116; 0.2214; 0.2301; 0.2378;\n 0.2448; 0.2513; 0.2574; 0.2633; 0.2691; 0.2747;\n 0.2801; 0.2852; 0.2899; 0.294; 0.2979; 0.3016;\n 0.3055; 0.3097; 0.3141; 0.3184; 0.3224; 0.3257;\n 0.3284; 0.3307; 0.3327; 0.3346; 0.3364; 0.3383;\n 0.3403; 0.3425; 0.3448; 0.3472; 0.3495; 0.3519;\n 0.3541; 0.3562; 0.3581; 0.3597; 0.3609; 0.3613;\n 0.3609; 0.3595; 0.3581; 0.3558; 0.3452; 0.3194;\n 0.2807; 0.2339; 0.1839; 0.1352; 0.0911; 0.0548;\n 0.0295; 0.0166; 0.0112; 0.0077; 0.005; 0.0032;\n 0.0021; 0.0015; 0.0012; 0.001; 0.0009; 0.0008;\n 0.0006; 0.0005; 0.0003; 0.0001; 0.0\n |]\n\n(* SDSS *)\n\nlet sdss_r_wave =\n [|\n 5380.0; 5405.0; 5430.0; 5455.0; 5480.0; 5505.0;\n 5530.0; 5555.0; 5580.0; 5605.0; 5630.0; 5655.0;\n 5680.0; 5705.0; 5730.0; 5755.0; 5780.0; 5805.0;\n 5830.0; 5855.0; 5880.0; 5905.0; 5930.0; 5955.0;\n 5980.0; 6005.0; 6030.0; 6055.0; 6080.0; 6105.0;\n 6130.0; 6155.0; 6180.0; 6205.0; 6230.0; 6255.0;\n 6280.0; 6305.0; 6330.0; 6355.0; 6380.0; 6405.0;\n 6430.0; 6455.0; 6480.0; 6505.0; 6530.0; 6555.0;\n 6580.0; 6605.0; 6630.0; 6655.0; 6680.0; 6705.0;\n 6730.0; 6755.0; 6780.0; 6805.0; 6830.0; 6855.0;\n 6880.0; 6905.0; 6930.0; 6955.0; 6980.0; 7005.0;\n 7030.0; 7055.0; 7080.0; 7105.0; 7130.0; 7155.0;\n 7180.0; 7205.0; 7230.0\n |]\n\nlet sdss_r_thru =\n [|\n 0.0; 0.0014; 0.0099; 0.0259; 0.0497; 0.0807;\n 0.1186; 0.1625; 0.2093; 0.2555; 0.2975; 0.3326;\n 0.3609; 0.3834; 0.401; 0.4147; 0.4253; 0.4333;\n 0.4395; 0.4446; 0.4489; 0.4527; 0.4563; 0.4599;\n 0.4634; 0.4665; 0.4689; 0.4703; 0.4711; 0.4717;\n 0.4727; 0.4744; 0.4767; 0.4792; 0.4819; 0.4844;\n 0.4867; 0.4887; 0.4902; 0.4909; 0.4912; 0.4912;\n 0.4912; 0.4914; 0.4915; 0.4912; 0.4901; 0.4878;\n 0.4852; 0.4818; 0.4697; 0.4421; 0.4009; 0.3499;\n 0.2924; 0.2318; 0.1715; 0.1152; 0.0687; 0.038;\n 0.0212; 0.0134; 0.0099; 0.0076; 0.0055; 0.0039;\n 0.0027; 0.002; 0.0015; 0.0012; 0.001; 0.0007;\n 0.0004; 0.0002; 0.0\n |]\n\n(* SDSS *)\n\nlet sdss_i_wave =\n [|\n 6430.0; 6455.0; 6480.0; 6505.0; 6530.0; 6555.0;\n 6580.0; 6605.0; 6630.0; 6655.0; 6680.0; 6705.0;\n 6730.0; 6755.0; 6780.0; 6805.0; 6830.0; 6855.0;\n 6880.0; 6905.0; 6930.0; 6955.0; 6980.0; 7005.0;\n 7030.0; 7055.0; 7080.0; 7105.0; 7130.0; 7155.0;\n 7180.0; 7205.0; 7230.0; 7255.0; 7280.0; 7305.0;\n 7330.0; 7355.0; 7380.0; 7405.0; 7430.0; 7455.0;\n 7480.0; 7505.0; 7530.0; 7555.0; 7580.0; 7605.0;\n 7630.0; 7655.0; 7680.0; 7705.0; 7730.0; 7755.0;\n 7780.0; 7805.0; 7830.0; 7855.0; 7880.0; 7905.0;\n 7930.0; 7955.0; 7980.0; 8005.0; 8030.0; 8055.0;\n 8080.0; 8105.0; 8130.0; 8155.0; 8180.0; 8205.0;\n 8230.0; 8255.0; 8280.0; 8305.0; 8330.0; 8355.0;\n 8380.0; 8405.0; 8430.0; 8455.0; 8480.0; 8505.0;\n 8530.0; 8555.0; 8580.0; 8605.0; 8630.0\n |]\n\nlet sdss_i_thru =\n [|\n 0.0; 0.0001; 0.0003; 0.0004; 0.0004; 0.0003;\n 0.0003; 0.0004; 0.0009; 0.0019; 0.0034; 0.0056;\n 0.0103; 0.0194; 0.0344; 0.0561; 0.0839; 0.1164;\n 0.1528; 0.1948; 0.2408; 0.2857; 0.3233; 0.3503;\n 0.3759; 0.399; 0.4162; 0.4233; 0.4165; 0.3943;\n 0.376; 0.3823; 0.3918; 0.3892; 0.3828; 0.382;\n 0.3884; 0.3872; 0.3821; 0.3787; 0.3759; 0.3727;\n 0.3681; 0.3618; 0.3565; 0.3554; 0.3478; 0.1473;\n 0.2096; 0.2648; 0.33; 0.3256; 0.3223; 0.3179;\n 0.3129; 0.3077; 0.3026; 0.298; 0.2944; 0.2921;\n 0.2916; 0.2921; 0.2927; 0.2923; 0.2896; 0.284;\n 0.2758; 0.2642; 0.2427; 0.2091; 0.1689; 0.1276;\n 0.0901; 0.0603; 0.0378; 0.0218; 0.0117; 0.0068;\n 0.0048; 0.0033; 0.002; 0.0013; 0.001; 0.0009;\n 0.0009; 0.0008; 0.0005; 0.0002; 0.0\n |]\n\n(* SDSS *)\n\nlet sdss_z_wave =\n [|\n 7730.0; 7755.0; 7780.0; 7805.0; 7830.0; 7855.0;\n 7880.0; 7905.0; 7930.0; 7955.0; 7980.0; 8005.0;\n 8030.0; 8055.0; 8080.0; 8105.0; 8130.0; 8155.0;\n 8180.0; 8205.0; 8230.0; 8255.0; 8280.0; 8305.0;\n 8330.0; 8355.0; 8380.0; 8405.0; 8430.0; 8455.0;\n 8480.0; 8505.0; 8530.0; 8555.0; 8580.0; 8605.0;\n 8630.0; 8655.0; 8680.0; 8705.0; 8730.0; 8755.0;\n 8780.0; 8805.0; 8830.0; 8855.0; 8880.0; 8905.0;\n 8930.0; 8955.0; 8980.0; 9005.0; 9030.0; 9055.0;\n 9080.0; 9105.0; 9130.0; 9155.0; 9180.0; 9205.0;\n 9230.0; 9255.0; 9280.0; 9305.0; 9330.0; 9355.0;\n 9380.0; 9405.0; 9430.0; 9455.0; 9480.0; 9505.0;\n 9530.0; 9555.0; 9580.0; 9605.0; 9630.0; 9655.0;\n 9680.0; 9705.0; 9730.0; 9755.0; 9780.0; 9805.0;\n 9830.0; 9855.0; 9880.0; 9905.0; 9930.0; 9955.0;\n 9980.0; 10005.0; 10030.0; 10055.0; 10080.0; 10105.0;\n 10130.0; 10155.0; 10180.0; 10205.0; 10230.0; 10255.0;\n 10280.0; 10305.0; 10330.0; 10355.0; 10380.0; 10405.0;\n 10430.0; 10455.0; 10480.0; 10505.0; 10530.0; 10555.0;\n 10580.0; 10605.0; 10630.0; 10655.0; 10680.0; 10705.0;\n 10730.0; 10755.0; 10780.0; 10805.0; 10830.0; 10855.0;\n 10880.0; 10905.0; 10930.0; 10955.0; 10980.0; 11005.0;\n 11030.0; 11055.0; 11080.0; 11105.0; 11130.0; 11155.0;\n 11180.0; 11205.0; 11230.0\n |]\n\nlet sdss_z_thru =\n [|\n 0.0; 0.0; 0.0001; 0.0001; 0.0001; 0.0002;\n 0.0002; 0.0003; 0.0005; 0.0007; 0.0011; 0.0017;\n 0.0027; 0.004; 0.0057; 0.0079; 0.0106; 0.0139;\n 0.0178; 0.0222; 0.0271; 0.0324; 0.0382; 0.0446;\n 0.0511; 0.0564; 0.0603; 0.0637; 0.0667; 0.0694;\n 0.0717; 0.0736; 0.0752; 0.0765; 0.0775; 0.0782;\n 0.0786; 0.0787; 0.0785; 0.078; 0.0772; 0.0763;\n 0.0751; 0.0738; 0.0723; 0.0708; 0.0693; 0.0674;\n 0.0632; 0.0581; 0.0543; 0.0526; 0.0523; 0.0522;\n 0.0512; 0.0496; 0.0481; 0.0473; 0.0476; 0.0482;\n 0.0476; 0.0447; 0.0391; 0.0329; 0.0283; 0.0264;\n 0.0271; 0.0283; 0.0275; 0.0254; 0.0252; 0.0256;\n 0.0246; 0.0244; 0.0252; 0.0258; 0.0265; 0.0274;\n 0.0279; 0.0271; 0.0252; 0.0236; 0.0227; 0.0222;\n 0.0216; 0.0208; 0.0196; 0.0183; 0.0171; 0.016;\n 0.0149; 0.0138; 0.0128; 0.0118; 0.0108; 0.0099;\n 0.0091; 0.0083; 0.0075; 0.0068; 0.0061; 0.0055;\n 0.005; 0.0045; 0.0041; 0.0037; 0.0033; 0.003;\n 0.0027; 0.0025; 0.0023; 0.0021; 0.0019; 0.0018;\n 0.0017; 0.0016; 0.0015; 0.0014; 0.0013; 0.0012;\n 0.0011; 0.001; 0.0009; 0.0008; 0.0008; 0.0007;\n 0.0006; 0.0006; 0.0006; 0.0005; 0.0005; 0.0004;\n 0.0004; 0.0003; 0.0003; 0.0002; 0.0002; 0.0001;\n 0.0001; 0.0; 0.0\n |]\n\n(* Johnson-Cousins *)\n\nlet johnson_u_wave =\n [|\n 3000.0; 3100.0; 3200.0; 3300.0; 3400.0; 3500.0;\n 3600.0; 3700.0; 3800.0; 3900.0; 4000.0; 4100.0;\n 4200.0\n |]\n\nlet johnson_u_thru =\n [|\n 0.0; 0.1; 0.61; 0.84; 0.93; 0.97;\n 1.0; 0.97; 0.73; 0.36; 0.05; 0.01;\n 0.0\n |]\n\n(* Johnson-Cousins *)\n\nlet johnson_b_wave =\n [|\n 3700.0; 3800.0; 4000.0; 4200.0; 4400.0; 4600.0;\n 4800.0; 5000.0; 5200.0; 5400.0; 5600.0\n |]\n\nlet johnson_b_thru =\n [|\n 0.0; 0.11; 0.92; 1.0; 0.94; 0.79;\n 0.58; 0.36; 0.15; 0.04; 0.0\n |]\n\n(* Johnson-Cousins *)\n\nlet johnson_v_wave =\n [|\n 4600.0; 4800.0; 5000.0; 5200.0; 5400.0; 5600.0;\n 5800.0; 6000.0; 6200.0; 6400.0; 6600.0; 6800.0;\n 7000.0; 7200.0; 7400.0\n |]\n\nlet johnson_v_thru =\n [|\n 0.0; 0.02; 0.38; 0.91; 0.98; 0.72;\n 0.62; 0.4; 0.2; 0.08; 0.02; 0.01;\n 0.01; 0.01; 0.0\n |]\n\n(* Johnson-Cousins *)\n\nlet cousins_r_wave =\n [|\n 5400.0; 5450.0; 5500.0; 5550.0; 5600.0; 5650.0;\n 5700.0; 5750.0; 5800.0; 5850.0; 5900.0; 5950.0;\n 6000.0; 6050.0; 6100.0; 6150.0; 6200.0; 6250.0;\n 6300.0; 6350.0; 6400.0; 6450.0; 6500.0; 6550.0;\n 6600.0; 6650.0; 6700.0; 6750.0; 6800.0; 6850.0;\n 6900.0; 6950.0; 7000.0; 7050.0; 7100.0; 7150.0;\n 7200.0; 7250.0; 7300.0; 7350.0; 7400.0; 7450.0;\n 7500.0; 7550.0; 7600.0; 7650.0; 7700.0; 7750.0;\n 7800.0; 7850.0; 7900.0; 7950.0; 8000.0\n |]\n\nlet cousins_r_thru =\n [|\n 0.0; 0.002; 0.01; 0.03; 0.07; 0.18;\n 0.4; 0.77; 0.89; 0.96; 0.99; 0.999;\n 1.0; 0.997; 0.99; 0.976; 0.96; 0.946;\n 0.93; 0.912; 0.895; 0.88; 0.86; 0.845;\n 0.825; 0.806; 0.788; 0.765; 0.742; 0.72;\n 0.7; 0.676; 0.65; 0.626; 0.6; 0.568;\n 0.53; 0.48; 0.395; 0.3; 0.215; 0.155;\n 0.12; 0.1; 0.085; 0.075; 0.06; 0.05;\n 0.04; 0.029; 0.02; 0.01; 0.0\n |]\n\n(* Johnson-Cousins *)\n\nlet cousins_i_wave =\n [|\n 7000.0; 7050.0; 7100.0; 7150.0; 7200.0; 7250.0;\n 7300.0; 7350.0; 7400.0; 7450.0; 7500.0; 7550.0;\n 7600.0; 7650.0; 7700.0; 7750.0; 7800.0; 7850.0;\n 7900.0; 7950.0; 8000.0; 8050.0; 8100.0; 8150.0;\n 8200.0; 8250.0; 8300.0; 8350.0; 8400.0; 8450.0;\n 8500.0; 8550.0; 8600.0; 8650.0; 8700.0; 8750.0;\n 8800.0; 8850.0; 8900.0; 8950.0; 9000.0; 9050.0;\n 9100.0\n |]\n\nlet cousins_i_thru =\n [|\n 0.0; 0.005; 0.02; 0.05; 0.1; 0.17;\n 0.33; 0.7; 0.82; 0.9; 0.95; 0.98;\n 0.99; 0.994; 0.98; 0.95; 0.913; 0.87;\n 0.83; 0.79; 0.75; 0.71; 0.673; 0.65;\n 0.63; 0.61; 0.58; 0.55; 0.51; 0.47;\n 0.405; 0.33; 0.25; 0.18; 0.14; 0.11;\n 0.08; 0.06; 0.035; 0.02; 0.01; 0.005;\n 0.0\n |]\n\n(* 2MASS *)\n\nlet twomass_j_wave =\n [|\n 10620.0; 10660.0; 10700.0; 10750.0; 10780.0; 10820.0;\n 10840.0; 10870.0; 10890.0; 10930.0; 10960.0; 11020.0;\n 11050.0; 11070.0; 11090.0; 11120.0; 11160.0; 11170.0;\n 11200.0; 11230.0; 11280.0; 11290.0; 11320.0; 11340.0;\n 11380.0; 11400.0; 11430.0; 11470.0; 11540.0; 11590.0;\n 11640.0; 11670.0; 11700.0; 11730.0; 11750.0; 11790.0;\n 11820.0; 11860.0; 11880.0; 11920.0; 11950.0; 11990.0;\n 12020.0; 12090.0; 12160.0; 12210.0; 12270.0; 12310.0;\n 12360.0; 12400.0; 12440.0; 12470.0; 12530.0; 12550.0;\n 12580.0; 12600.0; 12650.0; 12700.0; 12750.0; 12790.0;\n 12860.0; 12920.0; 12970.0; 13020.0; 13050.0; 13070.0;\n 13100.0; 13130.0; 13160.0; 13190.0; 13230.0; 13260.0;\n 13300.0; 13330.0; 13340.0; 13360.0; 13390.0; 13430.0;\n 13460.0; 13490.0; 13530.0; 13550.0; 13600.0; 13630.0;\n 13700.0; 13730.0; 13770.0; 13830.0; 13880.0; 13920.0;\n 13950.0; 13960.0; 13970.0; 13980.0; 14000.0; 14010.0;\n 14020.0; 14040.0; 14060.0; 14070.0; 14100.0; 14120.0;\n 14160.0; 14210.0; 14260.0; 14420.0; 14500.0\n |]\n\nlet twomass_j_thru =\n [|\n 0.0; 0.0004; 0.0015; 0.0027; 0.0055; 0.0123;\n 0.0203; 0.0306; 0.0405; 0.0515; 0.0564; 0.0718;\n 0.2736; 0.341; 0.3584; 0.3801; 0.3307; 0.2395;\n 0.2501; 0.2833; 0.2582; 0.2515; 0.5381; 0.2232;\n 0.5369; 0.1102; 0.5292; 0.2619; 0.3202; 0.1743;\n 0.607; 0.6179; 0.6763; 0.7279; 0.7465; 0.8304;\n 0.7903; 0.8096; 0.8369; 0.836; 0.7499; 0.708;\n 0.6988; 0.7049; 0.7004; 0.7328; 0.7057; 0.8424;\n 0.9219; 0.9525; 0.9676; 0.9595; 0.9227; 0.893;\n 0.8529; 0.8023; 0.7501; 0.6781; 0.6524; 0.6388;\n 0.6424; 0.6486; 0.6824; 0.7529; 0.7759; 0.8118;\n 0.777; 0.721; 0.9525; 0.8551; 0.8414; 1.0;\n 0.8947; 0.8549; 0.5379; 0.2799; 0.9065; 0.6893;\n 0.5533; 0.2432; 0.0144; 0.0002; 0.0401; 0.0045;\n 0.0003; 0.0372; 0.0005; 0.0; 0.0001; 0.0033;\n 0.0003; 0.0085; 0.0254; 0.1184; 0.0001; 0.0001;\n 0.0521; 0.0104; 0.0478; 0.0004; 0.0024; 0.0053;\n 0.0086; 0.0007; 0.0003; 0.0004; 0.0\n |]\n\n(* 2MASS *)\n\nlet twomass_h_wave =\n [|\n 12890.0; 13150.0; 13410.0; 13680.0; 13970.0; 14180.0;\n 14400.0; 14620.0; 14780.0; 14860.0; 14930.0; 15040.0;\n 15150.0; 15280.0; 15390.0; 15460.0; 15510.0; 15560.0;\n 15650.0; 15720.0; 15770.0; 15830.0; 15920.0; 15970.0;\n 16020.0; 16130.0; 16190.0; 16280.0; 16330.0; 16420.0;\n 16480.0; 16570.0; 16590.0; 16710.0; 16840.0; 17010.0;\n 17150.0; 17270.0; 17390.0; 17460.0; 17510.0; 17530.0;\n 17560.0; 17640.0; 17750.0; 17850.0; 17900.0; 17960.0;\n 18030.0; 18100.0; 18130.0; 18180.0; 18280.0; 18350.0;\n 18500.0; 18710.0; 18930.0; 19140.0\n |]\n\nlet twomass_h_thru =\n [|\n 0.0; 0.0; 0.0; 0.0; 0.0; 0.0;\n 0.0005; 0.028; 0.081; 0.287; 0.871; 0.201;\n 0.438; 0.686; 0.818; 0.882; 0.912; 0.927;\n 0.929; 0.873; 0.857; 0.883; 0.918; 0.927;\n 0.908; 0.926; 0.92; 0.924; 0.924; 0.942;\n 0.949; 0.981; 0.994; 1.0; 0.956; 0.924;\n 0.982; 0.992; 0.989; 0.979; 0.968; 0.937;\n 0.919; 0.842; 0.667; 0.269; 0.452; 0.173;\n 0.108; 0.071; 0.005; 0.02; 0.0004; 0.0;\n 0.0001; 0.0; 0.0; 0.0\n |]\n\n(* 2MASS *)\n\nlet twomass_ks_wave =\n [|\n 19000.0; 19150.0; 19270.0; 19340.0; 19390.0; 19480.0;\n 19570.0; 19620.0; 19690.0; 19760.0; 19810.0; 19890.0;\n 19900.0; 19980.0; 20080.0; 20140.0; 20190.0; 20280.0;\n 20370.0; 20450.0; 20610.0; 20720.0; 20750.0; 20820.0;\n 20890.0; 20990.0; 21060.0; 21130.0; 21200.0; 21240.0;\n 21380.0; 21450.0; 21550.0; 21690.0; 21760.0; 21850.0;\n 21970.0; 22080.0; 22130.0; 22180.0; 22320.0; 22370.0;\n 22480.0; 22560.0; 22600.0; 22630.0; 22650.0; 22700.0;\n 22720.0; 22760.0; 22770.0; 22810.0; 22840.0; 22860.0;\n 22910.0; 22930.0; 22950.0; 22970.0; 22990.0; 23060.0;\n 23110.0; 23160.0; 23200.0; 23250.0; 23280.0; 23350.0;\n 23390.0; 23440.0; 23460.0; 23520.0; 23610.0; 23630.0;\n 23700.0; 23750.0; 23840.0; 23990.0\n |]\n\nlet twomass_ks_thru =\n [|\n 0.0; 0.0; 0.0; 0.0002; 0.0005; 0.0054;\n 0.0119; 0.0197; 0.0422; 0.0873; 0.1528; 0.2482;\n 0.1902; 0.2339; 0.2946; 0.3982; 0.3366; 0.6207;\n 0.765; 0.7464; 0.6251; 0.7255; 0.6895; 0.7879;\n 0.8181; 0.8228; 0.8633; 0.8778; 0.8549; 0.8953;\n 0.9189; 0.9268; 0.9267; 0.9009; 0.9228; 0.8428;\n 0.9459; 0.9804; 0.9879; 0.9848; 0.9647; 0.9816;\n 0.9834; 0.9613; 0.9792; 1.0; 0.9632; 0.9812;\n 0.9681; 0.9109; 0.9821; 0.8896; 0.8918; 0.9424;\n 0.8404; 0.8042; 0.7077; 0.6576; 0.5607; 0.4437;\n 0.3482; 0.2302; 0.1626; 0.136; 0.0921; 0.0624;\n 0.0431; 0.034; 0.031; 0.0118; 0.0068; 0.0007;\n 0.003; 0.0021; 0.0004; 0.0\n |]\n\n(* Gaia DR3 *)\n\nlet gaia_g_wave =\n [|\n 3200.0; 3300.0; 3400.0; 3500.0; 3600.0; 3700.0;\n 3800.0; 3900.0; 4000.0; 4100.0; 4200.0; 4300.0;\n 4400.0; 4500.0; 4600.0; 4700.0; 4800.0; 4900.0;\n 5000.0; 5100.0; 5200.0; 5300.0; 5400.0; 5500.0;\n 5600.0; 5700.0; 5800.0; 5900.0; 6000.0; 6100.0;\n 6200.0; 6300.0; 6400.0; 6500.0; 6600.0; 6700.0;\n 6800.0; 6900.0; 7000.0; 7100.0; 7200.0; 7300.0;\n 7400.0; 7500.0; 7600.0; 7700.0; 7800.0; 7900.0;\n 8000.0; 8100.0; 8200.0; 8300.0; 8400.0; 8500.0;\n 8600.0; 8700.0; 8800.0; 8900.0; 9000.0; 9100.0;\n 9200.0; 9300.0; 9400.0; 9500.0; 9600.0; 9700.0;\n 9800.0; 9900.0; 10000.0; 10100.0; 10200.0; 10300.0;\n 10400.0; 10500.0\n |]\n\nlet gaia_g_thru =\n [|\n 2.37366962e-08; 0.00976875472; 0.0868837415; 0.125910068; 0.121442511; 0.109349045;\n 0.116293195; 0.204618287; 0.34084777; 0.433235889; 0.492915186; 0.532506055;\n 0.560121042; 0.58187167; 0.598921356; 0.612743401; 0.624456273; 0.634592054;\n 0.642876868; 0.651384274; 0.659234285; 0.665180853; 0.672624175; 0.677892686;\n 0.68337283; 0.688218588; 0.692909244; 0.698360314; 0.701281364; 0.705926392;\n 0.709945761; 0.712286557; 0.714900215; 0.716852196; 0.718062023; 0.717424017;\n 0.716404699; 0.713025742; 0.709495858; 0.702344476; 0.694885081; 0.682863231;\n 0.670880823; 0.654375536; 0.636105955; 0.615501457; 0.592399976; 0.567402553;\n 0.539583616; 0.510092228; 0.4791254; 0.447393833; 0.414784905; 0.38035191;\n 0.347263747; 0.313995072; 0.280491684; 0.249470941; 0.218314877; 0.189578109;\n 0.162072087; 0.137119296; 0.113758622; 0.0931891382; 0.074983285; 0.058819497;\n 0.0451523186; 0.0338677803; 0.0245381883; 0.0171045299; 0.0113958923; 0.00725157056;\n 0.00436700622; 0.00241251048\n |]\n\n(* Gaia DR3 *)\n\nlet gaia_bp_wave =\n [|\n 3250.0; 3300.0; 3350.0; 3400.0; 3450.0; 3500.0;\n 3550.0; 3600.0; 3650.0; 3700.0; 3750.0; 3800.0;\n 3850.0; 3900.0; 3950.0; 4000.0; 4050.0; 4100.0;\n 4150.0; 4200.0; 4250.0; 4300.0; 4350.0; 4400.0;\n 4450.0; 4500.0; 4550.0; 4600.0; 4650.0; 4700.0;\n 4750.0; 4800.0; 4850.0; 4900.0; 4950.0; 5000.0;\n 5050.0; 5100.0; 5150.0; 5200.0; 5250.0; 5300.0;\n 5350.0; 5400.0; 5450.0; 5500.0; 5550.0; 5600.0;\n 5650.0; 5700.0; 5750.0; 5800.0; 5850.0; 5900.0;\n 5950.0; 6000.0; 6050.0; 6100.0; 6150.0; 6200.0;\n 6250.0; 6300.0; 6350.0; 6400.0; 6450.0; 6500.0;\n 6550.0; 6600.0; 6650.0; 6700.0; 6750.0; 6800.0;\n 6850.0; 6900.0; 6950.0; 7000.0; 7050.0; 7100.0;\n 7150.0; 7200.0; 7250.0; 7300.0; 7350.0; 7400.0;\n 7450.0; 7500.0\n |]\n\nlet gaia_bp_thru =\n [|\n 3.87054116e-05; 0.0109458069; 0.0960352312; 0.209777042; 0.24623711; 0.184648618;\n 0.196988564; 0.235262373; 0.223965129; 0.204351616; 0.178318209; 0.162143918;\n 0.184059158; 0.254352193; 0.34761739; 0.432175816; 0.492469156; 0.533465853;\n 0.560569552; 0.578302699; 0.589904162; 0.59902208; 0.607578555; 0.615301491;\n 0.623213524; 0.626992584; 0.627863884; 0.627028071; 0.627574894; 0.629195435;\n 0.632206645; 0.634636782; 0.635341726; 0.635285854; 0.634064731; 0.631462795;\n 0.63078819; 0.630124067; 0.630179832; 0.630007723; 0.627664462; 0.623347947;\n 0.621221392; 0.620168751; 0.619955482; 0.622688637; 0.622951427; 0.619327372;\n 0.614279326; 0.608587274; 0.60526859; 0.613287749; 0.63076648; 0.643202692;\n 0.641217847; 0.627856079; 0.613625406; 0.613169161; 0.625579651; 0.649530382;\n 0.666534979; 0.666866457; 0.650929127; 0.620667611; 0.578699833; 0.528381533;\n 0.46236659; 0.341204564; 0.158484372; 0.0351559339; 0.00370522417; 0.000728910264;\n 0.000524017362; 0.000284946553; 0.000113610422; 2.0255692e-05; 6.08899493e-06; 4.25533546e-06;\n 2.88792024e-06; 9.89680713e-07; 4.24900727e-07; 1.27016225e-07; 1.16831386e-07; 7.93884518e-09;\n 1.27555036e-07; 2.17167412e-08\n |]\n\n(* Gaia DR3 *)\n\nlet gaia_rp_wave =\n [|\n 6100.0; 6150.0; 6200.0; 6250.0; 6300.0; 6350.0;\n 6400.0; 6450.0; 6500.0; 6550.0; 6600.0; 6650.0;\n 6700.0; 6750.0; 6800.0; 6850.0; 6900.0; 6950.0;\n 7000.0; 7050.0; 7100.0; 7150.0; 7200.0; 7250.0;\n 7300.0; 7350.0; 7400.0; 7450.0; 7500.0; 7550.0;\n 7600.0; 7650.0; 7700.0; 7750.0; 7800.0; 7850.0;\n 7900.0; 7950.0; 8000.0; 8050.0; 8100.0; 8150.0;\n 8200.0; 8250.0; 8300.0; 8350.0; 8400.0; 8450.0;\n 8500.0; 8550.0; 8600.0; 8650.0; 8700.0; 8750.0;\n 8800.0; 8850.0; 8900.0; 8950.0; 9000.0; 9050.0;\n 9100.0; 9150.0; 9200.0; 9250.0; 9300.0; 9350.0;\n 9400.0; 9450.0; 9500.0; 9550.0; 9600.0; 9650.0;\n 9700.0; 9750.0; 9800.0; 9850.0; 9900.0; 9950.0;\n 10000.0; 10050.0; 10100.0; 10150.0; 10200.0; 10250.0;\n 10300.0; 10350.0; 10400.0; 10450.0; 10500.0; 10550.0;\n 10600.0; 10650.0; 10700.0; 10750.0; 10800.0\n |]\n\nlet gaia_rp_thru =\n [|\n 0.0001067; 0.000705; 0.0089591; 0.0894186; 0.3945348; 0.6832151;\n 0.7284574; 0.6783742; 0.6932457; 0.6991653; 0.7068345; 0.7168661;\n 0.7258579; 0.7314582; 0.7317729; 0.729553; 0.7311262; 0.7341997;\n 0.7375911; 0.7377587; 0.7351913; 0.7317705; 0.7322348; 0.7341152;\n 0.7395558; 0.7439523; 0.7434368; 0.7401882; 0.7383857; 0.7400737;\n 0.7391916; 0.7378262; 0.7299905; 0.7234387; 0.7148353; 0.7081058;\n 0.7045418; 0.7029044; 0.703763; 0.7037788; 0.7012269; 0.698329;\n 0.6904644; 0.6830179; 0.6750185; 0.6668831; 0.6552453; 0.6437497;\n 0.6278626; 0.6142203; 0.5984866; 0.5817457; 0.5664293; 0.5505743;\n 0.5320554; 0.5156898; 0.4998404; 0.4817145; 0.4631831; 0.443315;\n 0.4236545; 0.4041978; 0.3837304; 0.3611222; 0.3409582; 0.320113;\n 0.2991975; 0.278412; 0.2555403; 0.2372075; 0.2165387; 0.1977315;\n 0.1787908; 0.1634732; 0.1453902; 0.1318572; 0.1142639; 0.0987333;\n 0.0815165; 0.0661173; 0.0521649; 0.0400458; 0.030169; 0.0228553;\n 0.0165918; 0.0122218; 0.0086189; 0.006114; 0.0042268; 0.0028113;\n 0.001905; 0.0012324; 0.0007693; 0.0004905; 0.0003028\n |]\n\n(* LSST/LSST.u — 60 points *)\nlet rubin_u_wave =\n [|\n 3.200000e+03; 3.215000e+03; 3.230000e+03; 3.245000e+03; 3.260000e+03; 3.275000e+03; 3.290000e+03; 3.305000e+03;\n 3.320000e+03; 3.335000e+03; 3.350000e+03; 3.365000e+03; 3.380000e+03; 3.395000e+03; 3.410000e+03; 3.425000e+03;\n 3.440000e+03; 3.455000e+03; 3.470000e+03; 3.485000e+03; 3.500000e+03; 3.515000e+03; 3.530000e+03; 3.545000e+03;\n 3.560000e+03; 3.575000e+03; 3.590000e+03; 3.605000e+03; 3.620000e+03; 3.635000e+03; 3.650000e+03; 3.665000e+03;\n 3.680000e+03; 3.695000e+03; 3.710000e+03; 3.725000e+03; 3.740000e+03; 3.755000e+03; 3.770000e+03; 3.785000e+03;\n 3.800000e+03; 3.815000e+03; 3.830000e+03; 3.845000e+03; 3.860000e+03; 3.875000e+03; 3.890000e+03; 3.905000e+03;\n 3.920000e+03; 3.935000e+03; 3.950000e+03; 3.965000e+03; 3.980000e+03; 3.995000e+03; 4.010000e+03; 4.025000e+03;\n 4.040000e+03; 4.055000e+03; 4.070000e+03; 4.085000e+03\n |]\n\nlet rubin_u_thru =\n [|\n 1.429550e-14; 5.824880e-03; 9.177360e-03; 1.413040e-02; 2.023590e-02; 2.751190e-02; 3.708220e-02; 4.640890e-02;\n 5.690710e-02; 6.560040e-02; 7.538320e-02; 8.192530e-02; 8.826960e-02; 9.514300e-02; 1.009060e-01; 1.072220e-01;\n 1.120190e-01; 1.179670e-01; 1.231450e-01; 1.283730e-01; 1.337000e-01; 1.381080e-01; 1.432610e-01; 1.478000e-01;\n 1.527230e-01; 1.573360e-01; 1.620670e-01; 1.666840e-01; 1.716940e-01; 1.764620e-01; 1.811790e-01; 1.858970e-01;\n 1.906280e-01; 1.950920e-01; 1.996840e-01; 2.041020e-01; 2.082430e-01; 2.126950e-01; 2.169940e-01; 2.214680e-01;\n 2.221790e-01; 2.194800e-01; 2.155940e-01; 2.047200e-01; 1.938620e-01; 1.822830e-01; 1.701760e-01; 1.575460e-01;\n 1.442870e-01; 1.308010e-01; 1.168670e-01; 1.024170e-01; 8.739400e-02; 7.196930e-02; 5.599790e-02; 3.968800e-02;\n 2.581460e-02; 1.741130e-02; 8.801300e-03; 2.931970e-04\n |]\n\n(* LSST/LSST.g — 60 points *)\nlet rubin_g_wave =\n [|\n 3.864000e+03; 3.894000e+03; 3.925000e+03; 3.955000e+03; 3.986000e+03; 4.016000e+03; 4.047000e+03; 4.078000e+03;\n 4.108000e+03; 4.139000e+03; 4.169000e+03; 4.200000e+03; 4.231000e+03; 4.261000e+03; 4.292000e+03; 4.322000e+03;\n 4.353000e+03; 4.384000e+03; 4.414000e+03; 4.445000e+03; 4.475000e+03; 4.506000e+03; 4.537000e+03; 4.567000e+03;\n 4.598000e+03; 4.628000e+03; 4.659000e+03; 4.690000e+03; 4.720000e+03; 4.751000e+03; 4.781000e+03; 4.812000e+03;\n 4.842000e+03; 4.873000e+03; 4.904000e+03; 4.934000e+03; 4.965000e+03; 4.995000e+03; 5.026000e+03; 5.057000e+03;\n 5.087000e+03; 5.118000e+03; 5.148000e+03; 5.179000e+03; 5.210000e+03; 5.240000e+03; 5.271000e+03; 5.301000e+03;\n 5.332000e+03; 5.363000e+03; 5.393000e+03; 5.424000e+03; 5.454000e+03; 5.485000e+03; 5.516000e+03; 5.546000e+03;\n 5.577000e+03; 5.607000e+03; 5.638000e+03; 5.669000e+03\n |]\n\nlet rubin_g_thru =\n [|\n 4.995720e-14; 1.504200e-02; 3.744180e-02; 7.157070e-02; 1.087810e-01; 1.464570e-01; 1.868340e-01; 2.288970e-01;\n 2.711150e-01; 3.057660e-01; 3.235230e-01; 3.295650e-01; 3.352140e-01; 3.406260e-01; 3.460290e-01; 3.509490e-01;\n 3.552980e-01; 3.595040e-01; 3.634800e-01; 3.669990e-01; 3.708900e-01; 3.741200e-01; 3.769130e-01; 3.795000e-01;\n 3.822150e-01; 3.843350e-01; 3.866750e-01; 3.889150e-01; 3.912650e-01; 3.926710e-01; 3.941770e-01; 3.949940e-01;\n 3.969120e-01; 3.983260e-01; 3.989640e-01; 3.998360e-01; 4.005050e-01; 4.012760e-01; 4.004640e-01; 4.001320e-01;\n 4.011100e-01; 4.024610e-01; 4.033200e-01; 4.036730e-01; 4.041100e-01; 4.038960e-01; 4.034070e-01; 4.035910e-01;\n 4.037410e-01; 4.053230e-01; 3.897960e-01; 3.553770e-01; 3.076690e-01; 2.585420e-01; 2.090320e-01; 1.607920e-01;\n 1.108210e-01; 6.215070e-02; 2.627740e-02; 8.320810e-04\n |]\n\n(* LSST/LSST.r — 60 points *)\nlet rubin_r_wave =\n [|\n 5.370000e+03; 5.398000e+03; 5.427000e+03; 5.455000e+03; 5.484000e+03; 5.513000e+03; 5.541000e+03; 5.570000e+03;\n 5.599000e+03; 5.627000e+03; 5.656000e+03; 5.684000e+03; 5.713000e+03; 5.742000e+03; 5.770000e+03; 5.799000e+03;\n 5.828000e+03; 5.856000e+03; 5.885000e+03; 5.913000e+03; 5.942000e+03; 5.971000e+03; 5.999000e+03; 6.028000e+03;\n 6.057000e+03; 6.085000e+03; 6.114000e+03; 6.142000e+03; 6.171000e+03; 6.200000e+03; 6.228000e+03; 6.257000e+03;\n 6.286000e+03; 6.314000e+03; 6.343000e+03; 6.371000e+03; 6.400000e+03; 6.429000e+03; 6.457000e+03; 6.486000e+03;\n 6.515000e+03; 6.543000e+03; 6.572000e+03; 6.600000e+03; 6.629000e+03; 6.658000e+03; 6.686000e+03; 6.715000e+03;\n 6.744000e+03; 6.772000e+03; 6.801000e+03; 6.829000e+03; 6.858000e+03; 6.887000e+03; 6.915000e+03; 6.944000e+03;\n 6.973000e+03; 7.001000e+03; 7.030000e+03; 7.059000e+03\n |]\n\nlet rubin_r_thru =\n [|\n 4.419110e-13; 2.309990e-02; 5.277260e-02; 9.905770e-02; 1.473500e-01; 1.958210e-01; 2.426460e-01; 2.913780e-01;\n 3.398970e-01; 3.865220e-01; 4.118140e-01; 4.177100e-01; 4.186640e-01; 4.200420e-01; 4.218920e-01; 4.241900e-01;\n 4.259640e-01; 4.276140e-01; 4.258520e-01; 4.267580e-01; 4.263350e-01; 4.286710e-01; 4.309470e-01; 4.325390e-01;\n 4.342740e-01; 4.365000e-01; 4.394080e-01; 4.418750e-01; 4.444110e-01; 4.462740e-01; 4.484890e-01; 4.503210e-01;\n 4.432760e-01; 4.520930e-01; 4.562140e-01; 4.584970e-01; 4.602860e-01; 4.625800e-01; 4.638100e-01; 4.629030e-01;\n 4.637690e-01; 4.655010e-01; 4.663710e-01; 4.690410e-01; 4.694020e-01; 4.697590e-01; 4.700940e-01; 4.709560e-01;\n 4.711620e-01; 4.671030e-01; 4.403810e-01; 3.889860e-01; 3.325260e-01; 2.460020e-01; 2.127620e-01; 1.675560e-01;\n 1.164150e-01; 6.332250e-02; 2.799150e-02; 9.340470e-04\n |]\n\n(* LSST/LSST.i — 60 points *)\nlet rubin_i_wave =\n [|\n 6.760000e+03; 6.786000e+03; 6.813000e+03; 6.839000e+03; 6.866000e+03; 6.892000e+03; 6.919000e+03; 6.946000e+03;\n 6.972000e+03; 6.999000e+03; 7.025000e+03; 7.052000e+03; 7.079000e+03; 7.105000e+03; 7.132000e+03; 7.158000e+03;\n 7.185000e+03; 7.212000e+03; 7.238000e+03; 7.265000e+03; 7.291000e+03; 7.318000e+03; 7.345000e+03; 7.371000e+03;\n 7.398000e+03; 7.424000e+03; 7.451000e+03; 7.478000e+03; 7.504000e+03; 7.531000e+03; 7.557000e+03; 7.584000e+03;\n 7.610000e+03; 7.637000e+03; 7.664000e+03; 7.690000e+03; 7.717000e+03; 7.743000e+03; 7.770000e+03; 7.797000e+03;\n 7.823000e+03; 7.850000e+03; 7.876000e+03; 7.903000e+03; 7.930000e+03; 7.956000e+03; 7.983000e+03; 8.009000e+03;\n 8.036000e+03; 8.063000e+03; 8.089000e+03; 8.116000e+03; 8.142000e+03; 8.169000e+03; 8.196000e+03; 8.222000e+03;\n 8.249000e+03; 8.275000e+03; 8.302000e+03; 8.329000e+03\n |]\n\nlet rubin_i_thru =\n [|\n 8.017680e-13; 2.428840e-02; 5.232930e-02; 1.008480e-01; 1.393260e-01; 1.741710e-01; 2.376070e-01; 2.935610e-01;\n 3.473580e-01; 3.942050e-01; 4.356630e-01; 4.633490e-01; 4.648150e-01; 4.656090e-01; 4.657820e-01; 4.634080e-01;\n 4.285660e-01; 4.468650e-01; 4.358040e-01; 4.434830e-01; 4.449560e-01; 4.513710e-01; 4.624340e-01; 4.619650e-01;\n 4.645030e-01; 4.657510e-01; 4.657690e-01; 4.658270e-01; 4.653330e-01; 4.654000e-01; 4.648360e-01; 4.636220e-01;\n 1.706640e-01; 2.698000e-01; 4.015270e-01; 4.520990e-01; 4.617890e-01; 4.617680e-01; 4.610620e-01; 4.606790e-01;\n 4.597540e-01; 4.585490e-01; 4.570360e-01; 4.535390e-01; 4.537330e-01; 4.537350e-01; 4.537880e-01; 4.513960e-01;\n 4.520760e-01; 4.306080e-01; 3.909260e-01; 3.398770e-01; 2.831320e-01; 2.292450e-01; 1.858430e-01; 1.434810e-01;\n 9.907340e-02; 5.212160e-02; 2.456530e-02; 8.945460e-04\n |]\n\n(* LSST/LSST.z — 60 points *)\nlet rubin_z_wave =\n [|\n 8.030000e+03; 8.052000e+03; 8.075000e+03; 8.098000e+03; 8.121000e+03; 8.144000e+03; 8.167000e+03; 8.190000e+03;\n 8.213000e+03; 8.236000e+03; 8.259000e+03; 8.282000e+03; 8.305000e+03; 8.328000e+03; 8.351000e+03; 8.374000e+03;\n 8.397000e+03; 8.420000e+03; 8.443000e+03; 8.466000e+03; 8.489000e+03; 8.512000e+03; 8.535000e+03; 8.558000e+03;\n 8.581000e+03; 8.604000e+03; 8.627000e+03; 8.650000e+03; 8.673000e+03; 8.696000e+03; 8.718000e+03; 8.741000e+03;\n 8.764000e+03; 8.787000e+03; 8.810000e+03; 8.833000e+03; 8.856000e+03; 8.879000e+03; 8.902000e+03; 8.925000e+03;\n 8.948000e+03; 8.971000e+03; 8.994000e+03; 9.017000e+03; 9.040000e+03; 9.063000e+03; 9.086000e+03; 9.109000e+03;\n 9.132000e+03; 9.155000e+03; 9.178000e+03; 9.201000e+03; 9.224000e+03; 9.247000e+03; 9.270000e+03; 9.293000e+03;\n 9.316000e+03; 9.339000e+03; 9.362000e+03; 9.385000e+03\n |]\n\nlet rubin_z_thru =\n [|\n 1.039400e-12; 1.983520e-02; 4.060140e-02; 7.732710e-02; 1.178440e-01; 1.535950e-01; 1.866070e-01; 2.287220e-01;\n 2.798410e-01; 2.967030e-01; 3.561920e-01; 3.877460e-01; 4.214340e-01; 4.411140e-01; 4.470510e-01; 4.475610e-01;\n 4.480250e-01; 4.471650e-01; 4.471040e-01; 4.466870e-01; 4.464830e-01; 4.458890e-01; 4.466020e-01; 4.471110e-01;\n 4.474290e-01; 4.476080e-01; 4.474960e-01; 4.474640e-01; 4.472600e-01; 4.469260e-01; 4.456110e-01; 4.441420e-01;\n 4.428510e-01; 4.412100e-01; 4.406520e-01; 4.402080e-01; 4.389510e-01; 4.380910e-01; 4.371550e-01; 4.326220e-01;\n 4.198210e-01; 3.961630e-01; 3.715530e-01; 3.863600e-01; 4.102820e-01; 3.960420e-01; 3.749220e-01; 3.652600e-01;\n 3.317280e-01; 2.897530e-01; 2.623570e-01; 2.428380e-01; 2.063130e-01; 1.666310e-01; 1.330530e-01; 8.828810e-02;\n 4.608740e-02; 1.844720e-02; 9.801210e-03; 2.278930e-04\n |]\n\n(* LSST/LSST.y — 60 points *)\nlet rubin_y_wave =\n [|\n 9.084000e+03; 9.116000e+03; 9.148000e+03; 9.180000e+03; 9.213000e+03; 9.245000e+03; 9.277000e+03; 9.310000e+03;\n 9.342000e+03; 9.374000e+03; 9.406000e+03; 9.439000e+03; 9.471000e+03; 9.503000e+03; 9.536000e+03; 9.568000e+03;\n 9.600000e+03; 9.632000e+03; 9.665000e+03; 9.697000e+03; 9.729000e+03; 9.762000e+03; 9.794000e+03; 9.826000e+03;\n 9.858000e+03; 9.891000e+03; 9.923000e+03; 9.955000e+03; 9.988000e+03; 1.002000e+04; 1.005200e+04; 1.008400e+04;\n 1.011700e+04; 1.014900e+04; 1.018100e+04; 1.021400e+04; 1.024600e+04; 1.027800e+04; 1.031000e+04; 1.034300e+04;\n 1.037500e+04; 1.040700e+04; 1.044000e+04; 1.047200e+04; 1.050400e+04; 1.053600e+04; 1.056900e+04; 1.060100e+04;\n 1.063300e+04; 1.066600e+04; 1.069800e+04; 1.073000e+04; 1.076200e+04; 1.079500e+04; 1.082700e+04; 1.085900e+04;\n 1.089200e+04; 1.092400e+04; 1.095600e+04; 1.098900e+04\n |]\n\nlet rubin_y_thru =\n [|\n 4.969710e-13; 2.294700e-02; 5.618380e-02; 9.902450e-02; 1.553400e-01; 1.976760e-01; 2.337680e-01; 2.243520e-01;\n 1.759080e-01; 2.175850e-01; 2.643880e-01; 2.261180e-01; 2.467280e-01; 2.298940e-01; 2.452360e-01; 2.340200e-01;\n 2.363950e-01; 2.468170e-01; 2.376270e-01; 2.601730e-01; 2.440820e-01; 2.274890e-01; 2.249930e-01; 2.233450e-01;\n 2.183580e-01; 2.090180e-01; 1.985730e-01; 1.887810e-01; 1.784510e-01; 1.678880e-01; 1.585850e-01; 1.488040e-01;\n 1.392830e-01; 1.302160e-01; 1.212720e-01; 1.124510e-01; 1.041190e-01; 9.340160e-02; 8.306840e-02; 7.337710e-02;\n 6.467130e-02; 5.662540e-02; 4.896100e-02; 4.205990e-02; 3.581300e-02; 3.014650e-02; 2.491370e-02; 2.058860e-02;\n 1.681940e-02; 1.360580e-02; 1.116100e-02; 9.274090e-03; 7.773550e-03; 6.345540e-03; 5.132820e-03; 3.961830e-03;\n 3.044690e-03; 2.228830e-03; 1.600220e-03; 1.225340e-03\n |]\n\n(* Euclid/VIS.vis — 60 points *)\nlet euclid_vis_wave =\n [|\n 4.369190e+03; 4.459140e+03; 4.549090e+03; 4.639040e+03; 4.738980e+03; 4.828920e+03; 4.918870e+03; 5.018810e+03;\n 5.108760e+03; 5.198710e+03; 5.298650e+03; 5.388590e+03; 5.478540e+03; 5.578480e+03; 5.668430e+03; 5.758380e+03;\n 5.858320e+03; 5.948270e+03; 6.038210e+03; 6.138150e+03; 6.228100e+03; 6.318050e+03; 6.417990e+03; 6.507940e+03;\n 6.597880e+03; 6.697820e+03; 6.787770e+03; 6.877720e+03; 6.977660e+03; 7.067610e+03; 7.157550e+03; 7.247500e+03;\n 7.347440e+03; 7.437390e+03; 7.527340e+03; 7.627280e+03; 7.717230e+03; 7.807170e+03; 7.907110e+03; 7.997060e+03;\n 8.087010e+03; 8.186950e+03; 8.276900e+03; 8.366840e+03; 8.466780e+03; 8.556730e+03; 8.646680e+03; 8.746620e+03;\n 8.836570e+03; 8.926510e+03; 9.026460e+03; 9.116400e+03; 9.206350e+03; 9.306290e+03; 9.396240e+03; 9.486180e+03;\n 9.586130e+03; 9.676070e+03; 9.766020e+03; 9.865960e+03\n |]\n\nlet euclid_vis_thru =\n [|\n 5.667901e-04; 1.630730e-03; 4.172531e-03; 1.124922e-03; 2.177489e-03; 3.386911e-03; 3.641207e-03; 1.371951e-02;\n 9.284691e-03; 1.350449e-02; 2.210145e-02; 5.507258e-02; 7.012943e-01; 7.157499e-01; 7.257763e-01; 7.345625e-01;\n 7.426705e-01; 7.485474e-01; 7.527962e-01; 7.552509e-01; 7.566956e-01; 7.574142e-01; 7.585293e-01; 7.588235e-01;\n 7.574724e-01; 7.558434e-01; 7.571670e-01; 7.570556e-01; 7.567355e-01; 7.559193e-01; 7.545203e-01; 7.533978e-01;\n 7.508411e-01; 7.461720e-01; 7.414610e-01; 7.350350e-01; 7.301607e-01; 7.229322e-01; 7.122474e-01; 7.009171e-01;\n 6.870897e-01; 6.690291e-01; 6.520249e-01; 6.316748e-01; 6.043242e-01; 5.787996e-01; 5.515213e-01; 5.182027e-01;\n 4.851838e-01; 4.521600e-01; 4.136761e-01; 3.779477e-01; 3.006216e-01; 1.259014e-02; 1.804853e-03; 2.027216e-03;\n 1.289551e-03; 6.535986e-04; 7.194299e-04; 4.038814e-04\n |]\n\n(* Euclid/NISP.Y — 60 points *)\nlet euclid_y_wave =\n [|\n 9.330000e+03; 9.380000e+03; 9.430000e+03; 9.480000e+03; 9.540000e+03; 9.590000e+03; 9.640000e+03; 9.700000e+03;\n 9.750000e+03; 9.800000e+03; 9.850000e+03; 9.910000e+03; 9.960000e+03; 1.001000e+04; 1.007000e+04; 1.012000e+04;\n 1.017000e+04; 1.022000e+04; 1.028000e+04; 1.033000e+04; 1.038000e+04; 1.044000e+04; 1.049000e+04; 1.054000e+04;\n 1.059000e+04; 1.065000e+04; 1.070000e+04; 1.075000e+04; 1.081000e+04; 1.086000e+04; 1.091000e+04; 1.096000e+04;\n 1.102000e+04; 1.107000e+04; 1.112000e+04; 1.118000e+04; 1.123000e+04; 1.128000e+04; 1.133000e+04; 1.139000e+04;\n 1.144000e+04; 1.149000e+04; 1.155000e+04; 1.160000e+04; 1.165000e+04; 1.170000e+04; 1.176000e+04; 1.181000e+04;\n 1.186000e+04; 1.192000e+04; 1.197000e+04; 1.202000e+04; 1.207000e+04; 1.213000e+04; 1.218000e+04; 1.223000e+04;\n 1.229000e+04; 1.234000e+04; 1.239000e+04; 1.245000e+04\n |]\n\nlet euclid_y_thru =\n [|\n 1.401100e-04; 9.044250e-04; 2.786910e-02; 1.932270e-01; 7.417070e-01; 7.539270e-01; 7.683890e-01; 7.725430e-01;\n 7.736840e-01; 7.748860e-01; 7.762310e-01; 7.789970e-01; 7.770510e-01; 7.763010e-01; 7.815180e-01; 7.831650e-01;\n 7.782630e-01; 7.784520e-01; 7.745790e-01; 7.736720e-01; 7.806700e-01; 7.784520e-01; 7.792640e-01; 7.782380e-01;\n 7.713900e-01; 7.704340e-01; 7.682960e-01; 7.640030e-01; 7.705670e-01; 7.696290e-01; 7.648760e-01; 7.663780e-01;\n 7.621110e-01; 7.587620e-01; 7.627900e-01; 7.664810e-01; 7.633970e-01; 7.646770e-01; 7.645850e-01; 7.692010e-01;\n 7.718380e-01; 7.719360e-01; 7.718390e-01; 7.717390e-01; 7.688890e-01; 7.725960e-01; 7.771790e-01; 7.775880e-01;\n 7.796710e-01; 7.805190e-01; 7.824850e-01; 7.843440e-01; 7.759730e-01; 3.002440e-01; 7.116640e-02; 1.522080e-02;\n 3.885880e-03; 1.834220e-03; 1.115430e-03; 6.624430e-04\n |]\n\n(* Euclid/NISP.J — 60 points *)\nlet euclid_j_wave =\n [|\n 1.141000e+04; 1.148000e+04; 1.156000e+04; 1.164000e+04; 1.172000e+04; 1.180000e+04; 1.188000e+04; 1.196000e+04;\n 1.204000e+04; 1.212000e+04; 1.220000e+04; 1.228000e+04; 1.236000e+04; 1.244000e+04; 1.252000e+04; 1.260000e+04;\n 1.268000e+04; 1.276000e+04; 1.284000e+04; 1.292000e+04; 1.299000e+04; 1.307000e+04; 1.315000e+04; 1.323000e+04;\n 1.331000e+04; 1.339000e+04; 1.347000e+04; 1.355000e+04; 1.363000e+04; 1.371000e+04; 1.379000e+04; 1.387000e+04;\n 1.395000e+04; 1.403000e+04; 1.411000e+04; 1.419000e+04; 1.427000e+04; 1.435000e+04; 1.443000e+04; 1.451000e+04;\n 1.458000e+04; 1.466000e+04; 1.474000e+04; 1.482000e+04; 1.490000e+04; 1.498000e+04; 1.506000e+04; 1.514000e+04;\n 1.522000e+04; 1.530000e+04; 1.538000e+04; 1.546000e+04; 1.554000e+04; 1.562000e+04; 1.570000e+04; 1.578000e+04;\n 1.586000e+04; 1.594000e+04; 1.602000e+04; 1.610000e+04\n |]\n\nlet euclid_j_thru =\n [|\n 1.576900e-04; 4.015150e-04; 3.417080e-03; 1.226570e-01; 7.426110e-01; 7.817110e-01; 7.813510e-01; 7.840630e-01;\n 7.888400e-01; 7.907170e-01; 7.833700e-01; 7.884310e-01; 7.896350e-01; 7.852690e-01; 7.966270e-01; 7.958310e-01;\n 7.988340e-01; 7.953290e-01; 7.964360e-01; 7.932720e-01; 7.885410e-01; 7.955600e-01; 7.943190e-01; 7.956280e-01;\n 8.027170e-01; 8.039270e-01; 8.032210e-01; 7.995800e-01; 8.013920e-01; 8.024890e-01; 7.976110e-01; 7.968730e-01;\n 7.954540e-01; 7.861820e-01; 7.882250e-01; 7.912090e-01; 7.856070e-01; 7.868450e-01; 7.890830e-01; 7.843430e-01;\n 7.847770e-01; 7.870230e-01; 7.847020e-01; 7.808650e-01; 7.816630e-01; 7.821060e-01; 7.840170e-01; 7.832910e-01;\n 7.825870e-01; 7.872660e-01; 7.816920e-01; 7.772310e-01; 7.810290e-01; 6.745540e-01; 2.009210e-01; 2.169140e-02;\n 3.511730e-03; 9.199010e-04; 3.143400e-04; 1.453270e-04\n |]\n\n(* Euclid/NISP.H — 60 points *)\nlet euclid_h_wave =\n [|\n 1.480000e+04; 1.489000e+04; 1.499000e+04; 1.509000e+04; 1.519000e+04; 1.529000e+04; 1.539000e+04; 1.549000e+04;\n 1.559000e+04; 1.569000e+04; 1.579000e+04; 1.589000e+04; 1.599000e+04; 1.609000e+04; 1.619000e+04; 1.629000e+04;\n 1.639000e+04; 1.649000e+04; 1.659000e+04; 1.669000e+04; 1.678000e+04; 1.688000e+04; 1.698000e+04; 1.708000e+04;\n 1.718000e+04; 1.728000e+04; 1.738000e+04; 1.748000e+04; 1.758000e+04; 1.768000e+04; 1.778000e+04; 1.788000e+04;\n 1.798000e+04; 1.808000e+04; 1.818000e+04; 1.828000e+04; 1.838000e+04; 1.848000e+04; 1.858000e+04; 1.868000e+04;\n 1.877000e+04; 1.887000e+04; 1.897000e+04; 1.907000e+04; 1.917000e+04; 1.927000e+04; 1.937000e+04; 1.947000e+04;\n 1.957000e+04; 1.967000e+04; 1.977000e+04; 1.987000e+04; 1.997000e+04; 2.007000e+04; 2.017000e+04; 2.027000e+04;\n 2.037000e+04; 2.047000e+04; 2.057000e+04; 2.067000e+04\n |]\n\nlet euclid_h_thru =\n [|\n 1.433800e-04; 3.416300e-04; 1.371980e-03; 1.165150e-02; 2.166120e-01; 7.653910e-01; 7.770660e-01; 7.766830e-01;\n 7.792960e-01; 7.733530e-01; 7.817380e-01; 7.820990e-01; 7.830000e-01; 7.815810e-01; 7.808620e-01; 7.824440e-01;\n 7.788240e-01; 7.785320e-01; 7.810690e-01; 7.777990e-01; 7.773880e-01; 7.825950e-01; 7.841750e-01; 7.836200e-01;\n 7.841940e-01; 7.853450e-01; 7.824440e-01; 7.815110e-01; 7.833050e-01; 7.838540e-01; 7.843000e-01; 7.839500e-01;\n 7.850730e-01; 7.855760e-01; 7.873750e-01; 7.901270e-01; 7.888930e-01; 7.901500e-01; 7.918380e-01; 7.927560e-01;\n 7.897410e-01; 7.867750e-01; 7.847490e-01; 7.831360e-01; 7.785180e-01; 7.761400e-01; 7.759030e-01; 7.723980e-01;\n 7.665570e-01; 7.651970e-01; 7.639360e-01; 7.611570e-01; 7.567070e-01; 7.462570e-01; 5.895400e-01; 1.360190e-01;\n 1.785060e-02; 3.091850e-03; 7.171190e-04; 1.507920e-04\n |]\n"
},
{
"path": "dev/umbra/lib/filters.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet f64 = Nx.float64\nlet angstrom_to_m = 1e-10\n\nlet make wave_a thru_a =\n let n = Array.length wave_a in\n let w = Nx.create f64 [| n |] wave_a in\n let w = Nx.mul_s w angstrom_to_m in\n let t = Nx.create f64 [| n |] thru_a in\n Photometry.bandpass ~wavelength:(Unit.Length.of_tensor w) ~throughput:t\n\n(* SDSS *)\n\nlet sdss_u = make Filter_data.sdss_u_wave Filter_data.sdss_u_thru\nlet sdss_g = make Filter_data.sdss_g_wave Filter_data.sdss_g_thru\nlet sdss_r = make Filter_data.sdss_r_wave Filter_data.sdss_r_thru\nlet sdss_i = make Filter_data.sdss_i_wave Filter_data.sdss_i_thru\nlet sdss_z = make Filter_data.sdss_z_wave Filter_data.sdss_z_thru\n\n(* Johnson-Cousins *)\n\nlet johnson_u = make Filter_data.johnson_u_wave Filter_data.johnson_u_thru\nlet johnson_b = make Filter_data.johnson_b_wave Filter_data.johnson_b_thru\nlet johnson_v = make Filter_data.johnson_v_wave Filter_data.johnson_v_thru\nlet cousins_r = make Filter_data.cousins_r_wave Filter_data.cousins_r_thru\nlet cousins_i = make Filter_data.cousins_i_wave Filter_data.cousins_i_thru\n\n(* 2MASS *)\n\nlet twomass_j = make Filter_data.twomass_j_wave Filter_data.twomass_j_thru\nlet twomass_h = make Filter_data.twomass_h_wave Filter_data.twomass_h_thru\nlet twomass_ks = make Filter_data.twomass_ks_wave Filter_data.twomass_ks_thru\n\n(* Gaia DR3 *)\n\nlet gaia_g = make Filter_data.gaia_g_wave Filter_data.gaia_g_thru\nlet gaia_bp = make Filter_data.gaia_bp_wave Filter_data.gaia_bp_thru\nlet gaia_rp = make Filter_data.gaia_rp_wave Filter_data.gaia_rp_thru\n\n(* Rubin/LSST *)\n\nlet rubin_u = make Filter_data.rubin_u_wave Filter_data.rubin_u_thru\nlet rubin_g = make Filter_data.rubin_g_wave Filter_data.rubin_g_thru\nlet rubin_r = make Filter_data.rubin_r_wave Filter_data.rubin_r_thru\nlet rubin_i = make Filter_data.rubin_i_wave Filter_data.rubin_i_thru\nlet rubin_z = make Filter_data.rubin_z_wave Filter_data.rubin_z_thru\nlet rubin_y = make Filter_data.rubin_y_wave Filter_data.rubin_y_thru\n\n(* Euclid *)\n\nlet euclid_vis = make Filter_data.euclid_vis_wave Filter_data.euclid_vis_thru\nlet euclid_y = make Filter_data.euclid_y_wave Filter_data.euclid_y_thru\nlet euclid_j = make Filter_data.euclid_j_wave Filter_data.euclid_j_thru\nlet euclid_h = make Filter_data.euclid_h_wave Filter_data.euclid_h_thru\n"
},
{
"path": "dev/umbra/lib/filters.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Standard astronomical filter bandpasses.\n\n Tabulated transmission curves from the\n {{:https://svo2.cab.inta-csic.es/theory/fps/} SVO Filter Profile Service}.\n Each value is a pre-built {!Photometry.bandpass}.\n\n {[\n let mag = Photometry.ab_mag Filters.sdss_r sed\n ]} *)\n\n(** {1:sdss SDSS ugriz} *)\n\nval sdss_u : Photometry.bandpass\n(** [sdss_u] is the SDSS u-band (298--413 nm, 47 points). *)\n\nval sdss_g : Photometry.bandpass\n(** [sdss_g] is the SDSS g-band (363--583 nm, 89 points). *)\n\nval sdss_r : Photometry.bandpass\n(** [sdss_r] is the SDSS r-band (538--723 nm, 75 points). *)\n\nval sdss_i : Photometry.bandpass\n(** [sdss_i] is the SDSS i-band (643--863 nm, 89 points). *)\n\nval sdss_z : Photometry.bandpass\n(** [sdss_z] is the SDSS z-band (773--1123 nm, 141 points). *)\n\n(** {1:johnson Johnson-Cousins UBVRI} *)\n\nval johnson_u : Photometry.bandpass\n(** [johnson_u] is the Johnson U-band (300--420 nm, 13 points). *)\n\nval johnson_b : Photometry.bandpass\n(** [johnson_b] is the Johnson B-band (370--560 nm, 11 points). *)\n\nval johnson_v : Photometry.bandpass\n(** [johnson_v] is the Johnson V-band (460--740 nm, 15 points). *)\n\nval cousins_r : Photometry.bandpass\n(** [cousins_r] is the Cousins R-band (540--800 nm, 53 points). *)\n\nval cousins_i : Photometry.bandpass\n(** [cousins_i] is the Cousins I-band (700--910 nm, 43 points). *)\n\n(** {1:twomass 2MASS JHKs} *)\n\nval twomass_j : Photometry.bandpass\n(** [twomass_j] is the 2MASS J-band (1062--1450 nm, 107 points). *)\n\nval twomass_h : Photometry.bandpass\n(** [twomass_h] is the 2MASS H-band (1289--1914 nm, 58 points). *)\n\nval twomass_ks : Photometry.bandpass\n(** [twomass_ks] is the 2MASS Ks-band (1900--2399 nm, 76 points). *)\n\n(** {1:gaia Gaia DR3} *)\n\nval gaia_g : Photometry.bandpass\n(** [gaia_g] is the Gaia DR3 G-band (330--1040 nm, 74 points). *)\n\nval gaia_bp : Photometry.bandpass\n(** [gaia_bp] is the Gaia DR3 BP-band (328--748 nm, 86 points). *)\n\nval gaia_rp : Photometry.bandpass\n(** [gaia_rp] is the Gaia DR3 RP-band (618--1076 nm, 95 points). *)\n\n(** {1:rubin Rubin/LSST ugrizy} *)\n\nval rubin_u : Photometry.bandpass\n(** [rubin_u] is the Rubin/LSST u-band (320--409 nm, 60 points). *)\n\nval rubin_g : Photometry.bandpass\n(** [rubin_g] is the Rubin/LSST g-band (386--567 nm, 60 points). *)\n\nval rubin_r : Photometry.bandpass\n(** [rubin_r] is the Rubin/LSST r-band (537--706 nm, 60 points). *)\n\nval rubin_i : Photometry.bandpass\n(** [rubin_i] is the Rubin/LSST i-band (676--833 nm, 60 points). *)\n\nval rubin_z : Photometry.bandpass\n(** [rubin_z] is the Rubin/LSST z-band (803--935 nm, 60 points). *)\n\nval rubin_y : Photometry.bandpass\n(** [rubin_y] is the Rubin/LSST y-band (908--1099 nm, 60 points). *)\n\n(** {1:euclid Euclid} *)\n\nval euclid_vis : Photometry.bandpass\n(** [euclid_vis] is the Euclid VIS-band (437--987 nm, 60 points). *)\n\nval euclid_y : Photometry.bandpass\n(** [euclid_y] is the Euclid NISP Y-band (933--1245 nm, 60 points). *)\n\nval euclid_j : Photometry.bandpass\n(** [euclid_j] is the Euclid NISP J-band (1141--1610 nm, 60 points). *)\n\nval euclid_h : Photometry.bandpass\n(** [euclid_h] is the Euclid NISP H-band (1480--2067 nm, 60 points). *)\n"
},
{
"path": "dev/umbra/lib/fits/dune",
"content": "(library\n (name umbra_fits)\n (public_name umbra.fits)\n (private_modules fits_parser)\n (libraries nx nx.io talon unix))\n"
},
{
"path": "dev/umbra/lib/fits/fits_parser.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet err_truncated = \"Fits: unexpected end of file\"\nlet err_no_simple = \"Fits: missing SIMPLE keyword in primary HDU\"\nlet err_bad_tform msg = \"Fits: unsupported TFORM: \" ^ msg\nlet block_size = 2880\n\ntype keyword = { key : string; value : string; comment : string }\n\ntype header = {\n keywords : keyword list;\n xtension : string;\n bitpix : int;\n naxis : int array;\n data_bytes : int;\n}\n\ntype col_desc = {\n name : string;\n tform : char;\n repeat : int;\n width : int;\n tnull : int64 option;\n tscal : float;\n tzero : float;\n}\n\nlet swap16 buf pos =\n let b0 = Bytes.get_uint8 buf pos in\n let b1 = Bytes.get_uint8 buf (pos + 1) in\n Bytes.set_uint8 buf pos b1;\n Bytes.set_uint8 buf (pos + 1) b0\n\nlet swap32 buf pos =\n let b0 = Bytes.get_uint8 buf pos in\n let b1 = Bytes.get_uint8 buf (pos + 1) in\n let b2 = Bytes.get_uint8 buf (pos + 2) in\n let b3 = Bytes.get_uint8 buf (pos + 3) in\n Bytes.set_uint8 buf pos b3;\n Bytes.set_uint8 buf (pos + 1) b2;\n Bytes.set_uint8 buf (pos + 2) b1;\n Bytes.set_uint8 buf (pos + 3) b0\n\nlet swap64 buf pos =\n let b0 = Bytes.get_uint8 buf pos in\n let b1 = Bytes.get_uint8 buf (pos + 1) in\n let b2 = Bytes.get_uint8 buf (pos + 2) in\n let b3 = Bytes.get_uint8 buf (pos + 3) in\n let b4 = Bytes.get_uint8 buf (pos + 4) in\n let b5 = Bytes.get_uint8 buf (pos + 5) in\n let b6 = Bytes.get_uint8 buf (pos + 6) in\n let b7 = Bytes.get_uint8 buf (pos + 7) in\n Bytes.set_uint8 buf pos b7;\n Bytes.set_uint8 buf (pos + 1) b6;\n Bytes.set_uint8 buf (pos + 2) b5;\n Bytes.set_uint8 buf (pos + 3) b4;\n Bytes.set_uint8 buf (pos + 4) b3;\n Bytes.set_uint8 buf (pos + 5) b2;\n Bytes.set_uint8 buf (pos + 6) b1;\n Bytes.set_uint8 buf (pos + 7) b0\n\nlet trim_right s =\n let len = String.length s in\n let i = ref (len - 1) in\n while !i >= 0 && s.[!i] = ' ' do\n decr i\n done;\n if !i = len - 1 then s else String.sub s 0 (!i + 1)\n\nlet parse_card card =\n let key = trim_right (String.sub card 0 8) in\n if key = \"COMMENT\" || key = \"HISTORY\" then\n let content =\n if String.length card > 8 then\n trim_right (String.sub card 8 (String.length card - 8))\n else \"\"\n in\n { key; value = content; comment = \"\" }\n else if String.length card < 10 || card.[8] <> '=' || card.[9] <> ' ' then\n { key; value = \"\"; comment = \"\" }\n else\n let rest = String.sub card 10 (String.length card - 10) in\n let rest = String.trim rest in\n if String.length rest > 0 && rest.[0] = '\\'' then begin\n let len = String.length rest in\n let i = ref 1 in\n let buf = Buffer.create 68 in\n while !i < len do\n if rest.[!i] = '\\'' then\n begin if !i + 1 < len && rest.[!i + 1] = '\\'' then begin\n Buffer.add_char buf '\\'';\n i := !i + 2\n end\n else i := len\n end\n else begin\n Buffer.add_char buf rest.[!i];\n i := !i + 1\n end\n done;\n { key; value = trim_right (Buffer.contents buf); comment = \"\" }\n end\n else\n begin match String.index_opt rest '/' with\n | Some i ->\n let value = trim_right (String.sub rest 0 i) in\n let comment =\n String.trim (String.sub rest (i + 1) (String.length rest - i - 1))\n in\n { key; value; comment }\n | None -> { key; value = trim_right rest; comment = \"\" }\n end\n\nlet read_one_header ic =\n let keywords = ref [] in\n let found_end = ref false in\n let card_buf = Bytes.create 80 in\n while not !found_end do\n let block = Bytes.create block_size in\n (match In_channel.really_input ic block 0 block_size with\n | None -> failwith err_truncated\n | Some () -> ());\n for card_i = 0 to 35 do\n if not !found_end then begin\n Bytes.blit block (card_i * 80) card_buf 0 80;\n let card = Bytes.to_string card_buf in\n let key = trim_right (String.sub card 0 8) in\n if key = \"END\" then found_end := true\n else if key <> \"\" then keywords := parse_card card :: !keywords\n end\n done\n done;\n List.rev !keywords\n\nlet find_keyword keywords key =\n match List.find_opt (fun kw -> kw.key = key) keywords with\n | Some kw -> Some kw.value\n | None -> None\n\nlet find_keyword_int keywords key =\n match find_keyword keywords key with\n | Some v -> Some (int_of_string (String.trim v))\n | None -> None\n\nlet find_keyword_exn keywords key =\n match find_keyword keywords key with\n | Some v -> v\n | None -> failwith (\"Fits: missing required keyword \" ^ key)\n\nlet find_keyword_int_exn keywords key =\n int_of_string (String.trim (find_keyword_exn keywords key))\n\nlet compute_data_bytes keywords =\n let bitpix = find_keyword_int_exn keywords \"BITPIX\" in\n let naxis_n = find_keyword_int_exn keywords \"NAXIS\" in\n if naxis_n = 0 then 0\n else begin\n let total = ref (abs bitpix / 8) in\n for i = 1 to naxis_n do\n let key = Printf.sprintf \"NAXIS%d\" i in\n total := !total * find_keyword_int_exn keywords key\n done;\n let pcount =\n match find_keyword_int keywords \"PCOUNT\" with Some v -> v | None -> 0\n in\n let gcount =\n match find_keyword_int keywords \"GCOUNT\" with Some v -> v | None -> 1\n in\n (!total + pcount) * gcount\n end\n\nlet build_header keywords =\n let bitpix = find_keyword_int_exn keywords \"BITPIX\" in\n let naxis_n = find_keyword_int_exn keywords \"NAXIS\" in\n let naxis =\n Array.init naxis_n (fun i ->\n find_keyword_int_exn keywords (Printf.sprintf \"NAXIS%d\" (i + 1)))\n in\n let xtension =\n match find_keyword keywords \"XTENSION\" with Some v -> v | None -> \"\"\n in\n let data_bytes = compute_data_bytes keywords in\n { keywords; xtension; bitpix; naxis; data_bytes }\n\nlet read_headers ic =\n In_channel.seek ic 0L;\n let headers = ref [] in\n let first = ref true in\n let continue = ref true in\n while !continue do\n let keywords = try Some (read_one_header ic) with Failure _ -> None in\n match keywords with\n | None -> continue := false\n | Some keywords ->\n if !first then begin\n first := false;\n match find_keyword keywords \"SIMPLE\" with\n | Some _ -> ()\n | None -> failwith err_no_simple\n end;\n let hdr = build_header keywords in\n headers := hdr :: !headers;\n let data_blocks =\n if hdr.data_bytes = 0 then 0\n else (hdr.data_bytes + block_size - 1) / block_size\n in\n In_channel.seek ic\n (Int64.add (In_channel.pos ic)\n (Int64.of_int (data_blocks * block_size)))\n done;\n List.rev !headers\n\nlet seek_to_data ic headers hdu =\n if hdu < 0 || hdu >= List.length headers then\n failwith\n (Printf.sprintf \"Fits: HDU %d out of range (file has %d)\" hdu\n (List.length headers));\n In_channel.seek ic 0L;\n for i = 0 to hdu do\n let h = List.nth headers i in\n let found_end = ref false in\n while not !found_end do\n let block = Bytes.create block_size in\n (match In_channel.really_input ic block 0 block_size with\n | None -> failwith err_truncated\n | Some () -> ());\n for card_i = 0 to 35 do\n if not !found_end then begin\n let key = trim_right (Bytes.sub_string block (card_i * 80) 8) in\n if key = \"END\" then found_end := true\n end\n done\n done;\n if i < hdu then begin\n let data_blocks =\n if h.data_bytes = 0 then 0\n else (h.data_bytes + block_size - 1) / block_size\n in\n In_channel.seek ic\n (Int64.add (In_channel.pos ic)\n (Int64.of_int (data_blocks * block_size)))\n end\n done;\n let h = List.nth headers hdu in\n h.data_bytes\n\nlet parse_tform s =\n let s = String.trim s in\n let len = String.length s in\n if len = 0 then failwith (err_bad_tform \"empty\");\n let i = ref 0 in\n while !i < len && s.[!i] >= '0' && s.[!i] <= '9' do\n incr i\n done;\n let repeat = if !i = 0 then 1 else int_of_string (String.sub s 0 !i) in\n if !i >= len then failwith (err_bad_tform s);\n let code = s.[!i] in\n let width =\n match code with\n | 'L' -> 1\n | 'B' -> 1\n | 'I' -> 2\n | 'J' -> 4\n | 'K' -> 8\n | 'E' -> 4\n | 'D' -> 8\n | 'A' -> 1\n | c -> failwith (err_bad_tform (String.make 1 c))\n in\n (code, repeat, width)\n\nlet parse_bintable_cols hdr =\n let keywords = hdr.keywords in\n let tfields = find_keyword_int_exn keywords \"TFIELDS\" in\n List.init tfields (fun i ->\n let col = i + 1 in\n let name =\n match find_keyword keywords (Printf.sprintf \"TTYPE%d\" col) with\n | Some v -> v\n | None -> Printf.sprintf \"col%d\" col\n in\n let tform_s = find_keyword_exn keywords (Printf.sprintf \"TFORM%d\" col) in\n let tform, repeat, width = parse_tform tform_s in\n let tnull =\n match find_keyword keywords (Printf.sprintf \"TNULL%d\" col) with\n | Some v -> Some (Int64.of_string (String.trim v))\n | None -> None\n in\n let tscal =\n match find_keyword keywords (Printf.sprintf \"TSCAL%d\" col) with\n | Some v -> float_of_string (String.trim v)\n | None -> 1.0\n in\n let tzero =\n match find_keyword keywords (Printf.sprintf \"TZERO%d\" col) with\n | Some v -> float_of_string (String.trim v)\n | None -> 0.0\n in\n { name; tform; repeat; width; tnull; tscal; tzero })\n"
},
{
"path": "dev/umbra/lib/fits/fits_parser.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(**/**)\n\n(** Internal FITS parser. *)\n\ntype keyword = { key : string; value : string; comment : string }\n\ntype header = {\n keywords : keyword list;\n xtension : string;\n bitpix : int;\n naxis : int array;\n data_bytes : int;\n}\n\ntype col_desc = {\n name : string;\n tform : char;\n repeat : int;\n width : int;\n tnull : int64 option;\n tscal : float;\n tzero : float;\n}\n\nval read_headers : In_channel.t -> header list\nval seek_to_data : In_channel.t -> header list -> int -> int\nval parse_bintable_cols : header -> col_desc list\nval find_keyword : keyword list -> string -> string option\nval find_keyword_int : keyword list -> string -> int option\nval trim_right : string -> string\nval block_size : int\nval swap16 : bytes -> int -> unit\nval swap32 : bytes -> int -> unit\nval swap64 : bytes -> int -> unit\n\n(**/**)\n"
},
{
"path": "dev/umbra/lib/fits/umbra_fits.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet err_not_bintable = \"Fits.read_table: HDU is not a BINTABLE\"\nlet err_not_image = \"Fits.read_image: HDU is not an image\"\nlet err_unsupported_bitpix n = Printf.sprintf \"Fits: unsupported BITPIX %d\" n\nlet err_truncated_data = \"Fits: unexpected end of file in data\"\n\ntype header_card = { key : string; value : string; comment : string }\ntype hdu_type = Primary | Image | Bintable | Ascii_table\n\ntype hdu_info = {\n index : int;\n hdu_type : hdu_type;\n dimensions : int array;\n num_rows : int option;\n num_cols : int option;\n}\n\nlet hdu_type_of_header i (hdr : Fits_parser.header) =\n match hdr.xtension with\n | \"\" -> if i = 0 then Primary else Image\n | \"BINTABLE\" -> Bintable\n | \"TABLE\" -> Ascii_table\n | \"IMAGE\" -> Image\n | _ -> Image\n\nlet read_input ic buf n =\n match In_channel.really_input ic buf 0 n with\n | None -> failwith err_truncated_data\n | Some () -> ()\n\nlet info path =\n let ic = In_channel.open_bin path in\n Fun.protect\n ~finally:(fun () -> In_channel.close ic)\n (fun () ->\n let headers = Fits_parser.read_headers ic in\n List.mapi\n (fun i (hdr : Fits_parser.header) ->\n let ht = hdu_type_of_header i hdr in\n let num_rows, num_cols =\n match ht with\n | Bintable | Ascii_table ->\n let nrows =\n if Array.length hdr.naxis >= 2 then Some hdr.naxis.(1)\n else None\n in\n let ncols =\n Fits_parser.find_keyword_int hdr.keywords \"TFIELDS\"\n in\n (nrows, ncols)\n | _ -> (None, None)\n in\n {\n index = i;\n hdu_type = ht;\n dimensions = hdr.naxis;\n num_rows;\n num_cols;\n })\n headers)\n\nlet header ?(hdu = 0) path =\n let ic = In_channel.open_bin path in\n Fun.protect\n ~finally:(fun () -> In_channel.close ic)\n (fun () ->\n let headers = Fits_parser.read_headers ic in\n if hdu < 0 || hdu >= List.length headers then\n failwith (Printf.sprintf \"Fits.header: HDU %d out of range\" hdu);\n let h = List.nth headers hdu in\n List.map\n (fun (kw : Fits_parser.keyword) ->\n { key = kw.key; value = kw.value; comment = kw.comment })\n h.keywords)\n\nlet read_table ?(hdu = 1) path =\n let ic = In_channel.open_bin path in\n Fun.protect\n ~finally:(fun () -> In_channel.close ic)\n (fun () ->\n let headers = Fits_parser.read_headers ic in\n if hdu < 0 || hdu >= List.length headers then\n failwith (Printf.sprintf \"Fits.read_table: HDU %d out of range\" hdu);\n let h = List.nth headers hdu in\n (match hdu_type_of_header hdu h with\n | Bintable -> ()\n | _ -> failwith err_not_bintable);\n let cols = Fits_parser.parse_bintable_cols h in\n let nrows = if Array.length h.naxis >= 2 then h.naxis.(1) else 0 in\n let row_bytes = h.naxis.(0) in\n let (_ : int) = Fits_parser.seek_to_data ic headers hdu in\n let row_buf = Bytes.create row_bytes in\n let col_info =\n List.map\n (fun (cd : Fits_parser.col_desc) ->\n let elem_bytes = cd.repeat * cd.width in\n (cd, Bytes.create (nrows * elem_bytes), elem_bytes))\n cols\n in\n let col_offsets =\n let off = ref 0 in\n List.map\n (fun (cd : Fits_parser.col_desc) ->\n let o = !off in\n off := !off + (cd.repeat * cd.width);\n o)\n cols\n in\n for row = 0 to nrows - 1 do\n read_input ic row_buf row_bytes;\n List.iter2\n (fun offset (_cd, buf, elem_bytes) ->\n Bytes.blit row_buf offset buf (row * elem_bytes) elem_bytes)\n col_offsets col_info\n done;\n let err_vector name repeat =\n failwith\n (Printf.sprintf \"Fits: vector column '%s' (repeat=%d) not supported\"\n name repeat)\n in\n let talon_cols =\n List.map\n (fun (cd, buf, _) ->\n let col =\n match cd.Fits_parser.tform with\n | 'E' ->\n if cd.repeat <> 1 then err_vector cd.name cd.repeat;\n Talon.Col.float32\n (Array.init nrows (fun i ->\n let pos = i * 4 in\n Fits_parser.swap32 buf pos;\n let v =\n Int32.float_of_bits (Bytes.get_int32_le buf pos)\n in\n if cd.tzero = 0.0 && cd.tscal = 1.0 then v\n else (v *. cd.tscal) +. cd.tzero))\n | 'D' ->\n if cd.repeat <> 1 then err_vector cd.name cd.repeat;\n Talon.Col.float64\n (Array.init nrows (fun i ->\n let pos = i * 8 in\n Fits_parser.swap64 buf pos;\n let v =\n Int64.float_of_bits (Bytes.get_int64_le buf pos)\n in\n if cd.tzero = 0.0 && cd.tscal = 1.0 then v\n else (v *. cd.tscal) +. cd.tzero))\n | 'J' ->\n if cd.repeat <> 1 then err_vector cd.name cd.repeat;\n Talon.Col.int32\n (Array.init nrows (fun i ->\n let pos = i * 4 in\n Fits_parser.swap32 buf pos;\n let v = Bytes.get_int32_le buf pos in\n if cd.tzero = 0.0 && cd.tscal = 1.0 then v\n else\n Int32.of_float\n ((Int32.to_float v *. cd.tscal) +. cd.tzero)))\n | 'K' ->\n if cd.repeat <> 1 then err_vector cd.name cd.repeat;\n Talon.Col.int64\n (Array.init nrows (fun i ->\n let pos = i * 8 in\n Fits_parser.swap64 buf pos;\n let v = Bytes.get_int64_le buf pos in\n if cd.tzero = 0.0 && cd.tscal = 1.0 then v\n else\n Int64.of_float\n ((Int64.to_float v *. cd.tscal) +. cd.tzero)))\n | 'I' ->\n if cd.repeat <> 1 then err_vector cd.name cd.repeat;\n Talon.Col.int32\n (Array.init nrows (fun i ->\n let pos = i * 2 in\n Fits_parser.swap16 buf pos;\n let v = Bytes.get_int16_le buf pos in\n if cd.tzero = 0.0 && cd.tscal = 1.0 then Int32.of_int v\n else\n Int32.of_float\n ((Float.of_int v *. cd.tscal) +. cd.tzero)))\n | 'B' ->\n if cd.repeat <> 1 then err_vector cd.name cd.repeat;\n Talon.Col.int32\n (Array.init nrows (fun i ->\n let v = Bytes.get_uint8 buf i in\n if cd.tzero = 0.0 && cd.tscal = 1.0 then Int32.of_int v\n else\n Int32.of_float\n ((Float.of_int v *. cd.tscal) +. cd.tzero)))\n | 'L' ->\n if cd.repeat <> 1 then err_vector cd.name cd.repeat;\n Talon.Col.bool\n (Array.init nrows (fun i ->\n let c = Bytes.get buf i in\n c = 'T' || c = '\\x01'))\n | 'A' ->\n Talon.Col.string\n (Array.init nrows (fun i ->\n Fits_parser.trim_right\n (Bytes.sub_string buf (i * cd.repeat) cd.repeat)))\n | c -> failwith (Printf.sprintf \"Fits: unsupported TFORM '%c'\" c)\n in\n (cd.name, col))\n col_info\n in\n Talon.create talon_cols)\n\nlet find_keyword_float keywords key =\n match Fits_parser.find_keyword keywords key with\n | Some v -> Some (float_of_string (String.trim v))\n | None -> None\n\nlet read_image ?(hdu = 0) path =\n let ic = In_channel.open_bin path in\n Fun.protect\n ~finally:(fun () -> In_channel.close ic)\n (fun () ->\n let headers = Fits_parser.read_headers ic in\n if hdu < 0 || hdu >= List.length headers then\n failwith (Printf.sprintf \"Fits.read_image: HDU %d out of range\" hdu);\n let h = List.nth headers hdu in\n (match hdu_type_of_header hdu h with\n | Primary | Image -> ()\n | _ -> failwith err_not_image);\n let bscale =\n match find_keyword_float h.keywords \"BSCALE\" with\n | Some v -> v\n | None -> 1.0\n in\n let bzero =\n match find_keyword_float h.keywords \"BZERO\" with\n | Some v -> v\n | None -> 0.0\n in\n let has_scaling = bscale <> 1.0 || bzero <> 0.0 in\n let (_ : int) = Fits_parser.seek_to_data ic headers hdu in\n let shape = Array.to_list h.naxis |> List.rev |> Array.of_list in\n let total = Array.fold_left ( * ) 1 shape in\n let apply_scaling raw =\n Nx.add_s (Nx.mul_s (Nx.astype Nx.float64 raw) bscale) bzero\n in\n match h.bitpix with\n | 8 ->\n let buf = Bytes.create total in\n read_input ic buf total;\n let raw =\n Nx.create Nx.uint8 shape\n (Array.init total (fun i -> Bytes.get_uint8 buf i))\n in\n if has_scaling then Nx_io.P (apply_scaling raw) else Nx_io.P raw\n | 16 ->\n let buf = Bytes.create (total * 2) in\n read_input ic buf (total * 2);\n let raw =\n Nx.create Nx.int16 shape\n (Array.init total (fun i ->\n let pos = i * 2 in\n Fits_parser.swap16 buf pos;\n Bytes.get_int16_le buf pos))\n in\n if has_scaling then Nx_io.P (apply_scaling raw) else Nx_io.P raw\n | 32 ->\n let buf = Bytes.create (total * 4) in\n read_input ic buf (total * 4);\n let raw =\n Nx.create Nx.int32 shape\n (Array.init total (fun i ->\n let pos = i * 4 in\n Fits_parser.swap32 buf pos;\n Bytes.get_int32_le buf pos))\n in\n if has_scaling then Nx_io.P (apply_scaling raw) else Nx_io.P raw\n | 64 ->\n let buf = Bytes.create (total * 8) in\n read_input ic buf (total * 8);\n let raw =\n Nx.create Nx.int64 shape\n (Array.init total (fun i ->\n let pos = i * 8 in\n Fits_parser.swap64 buf pos;\n Bytes.get_int64_le buf pos))\n in\n if has_scaling then Nx_io.P (apply_scaling raw) else Nx_io.P raw\n | -32 ->\n let buf = Bytes.create (total * 4) in\n read_input ic buf (total * 4);\n let raw =\n Nx.create Nx.float32 shape\n (Array.init total (fun i ->\n let pos = i * 4 in\n Fits_parser.swap32 buf pos;\n Int32.float_of_bits (Bytes.get_int32_le buf pos)))\n in\n if has_scaling then Nx_io.P (apply_scaling raw) else Nx_io.P raw\n | -64 ->\n let buf = Bytes.create (total * 8) in\n read_input ic buf (total * 8);\n let raw =\n Nx.create Nx.float64 shape\n (Array.init total (fun i ->\n let pos = i * 8 in\n Fits_parser.swap64 buf pos;\n Int64.float_of_bits (Bytes.get_int64_le buf pos)))\n in\n if has_scaling then Nx_io.P (apply_scaling raw) else Nx_io.P raw\n | n -> failwith (err_unsupported_bitpix n))\n\nlet pad_to_block oc written =\n let rem = written mod Fits_parser.block_size in\n if rem > 0 then\n output_string oc (String.make (Fits_parser.block_size - rem) '\\x00')\n\nlet write_card oc key value =\n let card = Bytes.make 80 ' ' in\n Bytes.blit_string key 0 card 0 (Int.min 8 (String.length key));\n Bytes.set card 8 '=';\n Bytes.set card 9 ' ';\n let v = String.trim value in\n Bytes.blit_string v 0 card 10 (Int.min 70 (String.length v));\n output_bytes oc card\n\nlet write_card_str oc key value =\n write_card oc key (Printf.sprintf \"'%-8s'\" value)\n\nlet write_card_int oc key value =\n write_card oc key (Printf.sprintf \"%20d\" value)\n\nlet write_end oc cards_written =\n let card = Bytes.make 80 ' ' in\n Bytes.blit_string \"END\" 0 card 0 3;\n output_bytes oc card;\n let total_cards = cards_written + 1 in\n let rem = total_cards * 80 mod Fits_parser.block_size in\n if rem > 0 then\n output_string oc (String.make (Fits_parser.block_size - rem) ' ')\n\nlet write_empty_primary oc =\n write_card oc \"SIMPLE\" \" T\";\n write_card_int oc \"BITPIX\" 8;\n write_card_int oc \"NAXIS\" 0;\n write_end oc 3\n\nlet write_image_typed (type a b) ?(overwrite = true) path (tensor : (a, b) Nx.t)\n =\n if (not overwrite) && Sys.file_exists path then\n failwith (\"Fits.write_image: file exists: \" ^ path);\n let oc = Out_channel.open_bin path in\n Fun.protect\n ~finally:(fun () -> Out_channel.close oc)\n (fun () ->\n let shape = Nx.shape tensor in\n let ndim = Array.length shape in\n let fits_shape = Array.init ndim (fun i -> shape.(ndim - 1 - i)) in\n let total = Nx.numel tensor in\n let dt = Nx.dtype_to_string (Nx.dtype tensor) in\n let bitpix, elem_bytes =\n match dt with\n | \"uint8\" -> (8, 1)\n | \"int16\" -> (16, 2)\n | \"int32\" -> (32, 4)\n | \"int64\" -> (64, 8)\n | \"float32\" -> (-32, 4)\n | \"float64\" -> (-64, 8)\n | s -> failwith (\"Fits.write_image: unsupported dtype \" ^ s)\n in\n write_card oc \"SIMPLE\" \" T\";\n write_card_int oc \"BITPIX\" bitpix;\n write_card_int oc \"NAXIS\" ndim;\n for i = 0 to ndim - 1 do\n write_card_int oc (Printf.sprintf \"NAXIS%d\" (i + 1)) fits_shape.(i)\n done;\n write_end oc (3 + ndim);\n let flat = Nx.reshape [| total |] tensor in\n let arr = Nx.to_array flat in\n let data_bytes = total * elem_bytes in\n let buf = Bytes.create data_bytes in\n (match dt with\n | \"uint8\" ->\n Array.iteri\n (fun i (v : a) -> Bytes.set_uint8 buf i (Obj.magic v : int))\n arr\n | \"int16\" ->\n Array.iteri\n (fun i (v : a) ->\n let pos = i * 2 in\n Bytes.set_int16_le buf pos (Obj.magic v : int);\n Fits_parser.swap16 buf pos)\n arr\n | \"int32\" ->\n Array.iteri\n (fun i (v : a) ->\n let pos = i * 4 in\n Bytes.set_int32_le buf pos (Obj.magic v : int32);\n Fits_parser.swap32 buf pos)\n arr\n | \"int64\" ->\n Array.iteri\n (fun i (v : a) ->\n let pos = i * 8 in\n Bytes.set_int64_le buf pos (Obj.magic v : int64);\n Fits_parser.swap64 buf pos)\n arr\n | \"float32\" ->\n Array.iteri\n (fun i (v : a) ->\n let pos = i * 4 in\n Bytes.set_int32_le buf pos\n (Int32.bits_of_float (Obj.magic v : float));\n Fits_parser.swap32 buf pos)\n arr\n | \"float64\" ->\n Array.iteri\n (fun i (v : a) ->\n let pos = i * 8 in\n Bytes.set_int64_le buf pos\n (Int64.bits_of_float (Obj.magic v : float));\n Fits_parser.swap64 buf pos)\n arr\n | _ -> assert false);\n output_bytes oc buf;\n pad_to_block oc data_bytes)\n\nlet write_image ?overwrite path tensor =\n write_image_typed ?overwrite path tensor\n\nlet write_table ?(overwrite = true) path df =\n if (not overwrite) && Sys.file_exists path then\n failwith (\"Fits.write_table: file exists: \" ^ path);\n let oc = Out_channel.open_bin path in\n Fun.protect\n ~finally:(fun () -> Out_channel.close oc)\n (fun () ->\n write_empty_primary oc;\n let col_names = Talon.column_names df in\n let nrows = Talon.num_rows df in\n let ncols = List.length col_names in\n let col_info =\n List.map\n (fun name ->\n let col = Talon.get_column_exn df name in\n match Talon.Col.dtype col with\n | `Float32 -> (name, col, \"1E\", 4)\n | `Float64 -> (name, col, \"1D\", 8)\n | `Int32 -> (name, col, \"1J\", 4)\n | `Int64 -> (name, col, \"1K\", 8)\n | `String -> (\n match Talon.to_string_array df name with\n | Some arr ->\n let maxlen =\n Array.fold_left\n (fun acc v ->\n match v with\n | Some s -> max acc (String.length s)\n | None -> acc)\n 1 arr\n in\n (name, col, Printf.sprintf \"%dA\" maxlen, maxlen)\n | None -> failwith \"Fits.write_table: string column missing\")\n | `Bool -> (name, col, \"1L\", 1)\n | `Other -> failwith \"Fits.write_table: unsupported dtype\")\n col_names\n in\n let row_bytes =\n List.fold_left (fun acc (_, _, _, eb) -> acc + eb) 0 col_info\n in\n write_card_str oc \"XTENSION\" \"BINTABLE\";\n write_card_int oc \"BITPIX\" 8;\n write_card_int oc \"NAXIS\" 2;\n write_card_int oc \"NAXIS1\" row_bytes;\n write_card_int oc \"NAXIS2\" nrows;\n write_card_int oc \"PCOUNT\" 0;\n write_card_int oc \"GCOUNT\" 1;\n write_card_int oc \"TFIELDS\" ncols;\n let cards = ref 8 in\n List.iteri\n (fun i (name, _col, tform, _eb) ->\n let n = i + 1 in\n write_card_str oc (Printf.sprintf \"TTYPE%d\" n) name;\n write_card_str oc (Printf.sprintf \"TFORM%d\" n) tform;\n cards := !cards + 2)\n col_info;\n write_end oc !cards;\n let col_arrays =\n List.map\n (fun (name, col, _tform, _eb) ->\n match Talon.Col.dtype col with\n | `Float32 -> (\n match Talon.to_array Nx.float32 df name with\n | Some a -> `F32 a\n | None -> assert false)\n | `Float64 -> (\n match Talon.to_array Nx.float64 df name with\n | Some a -> `F64 a\n | None -> assert false)\n | `Int32 -> (\n match Talon.to_array Nx.int32 df name with\n | Some a -> `I32 a\n | None -> assert false)\n | `Int64 -> (\n match Talon.to_array Nx.int64 df name with\n | Some a -> `I64 a\n | None -> assert false)\n | `String -> (\n match Talon.to_string_array df name with\n | Some a -> `Str a\n | None -> assert false)\n | `Bool -> (\n match Talon.to_bool_array df name with\n | Some a -> `Bool a\n | None -> assert false)\n | `Other -> failwith \"Fits.write_table: unsupported dtype\")\n col_info\n in\n let row_buf = Bytes.create row_bytes in\n for row = 0 to nrows - 1 do\n let off = ref 0 in\n List.iter2\n (fun (_, _, _, eb) col_arr ->\n (match col_arr with\n | `F32 arr ->\n Bytes.set_int32_le row_buf !off (Int32.bits_of_float arr.(row));\n Fits_parser.swap32 row_buf !off\n | `F64 arr ->\n Bytes.set_int64_le row_buf !off (Int64.bits_of_float arr.(row));\n Fits_parser.swap64 row_buf !off\n | `I32 arr ->\n Bytes.set_int32_le row_buf !off arr.(row);\n Fits_parser.swap32 row_buf !off\n | `I64 arr ->\n Bytes.set_int64_le row_buf !off arr.(row);\n Fits_parser.swap64 row_buf !off\n | `Str arr -> (\n Bytes.fill row_buf !off eb ' ';\n match arr.(row) with\n | Some s ->\n let len = Int.min eb (String.length s) in\n Bytes.blit_string s 0 row_buf !off len\n | None -> ())\n | `Bool arr ->\n let v = match arr.(row) with Some true -> 'T' | _ -> 'F' in\n Bytes.set row_buf !off v);\n off := !off + eb)\n col_info col_arrays;\n output_bytes oc row_buf\n done;\n pad_to_block oc (nrows * row_bytes))\n"
},
{
"path": "dev/umbra/lib/fits/umbra_fits.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** FITS file I/O.\n\n Reads and writes {{:https://fits.gsfc.nasa.gov/fits_standard.html}FITS}\n files. Binary tables are loaded into {!Talon.t} dataframes and images into\n {!Nx.t} tensors. All data is converted from FITS big-endian on read and\n written as big-endian on write. *)\n\n(** {1:inspect Inspection} *)\n\n(** The type for FITS header data unit kinds. *)\ntype hdu_type =\n | Primary (** Primary HDU. *)\n | Image (** Image extension. *)\n | Bintable (** Binary table extension. *)\n | Ascii_table (** ASCII table extension. *)\n\ntype hdu_info = {\n index : int; (** Zero-based HDU index. *)\n hdu_type : hdu_type; (** Kind of HDU. *)\n dimensions : int array; (** NAXIS values. *)\n num_rows : int option; (** Row count for table HDUs. *)\n num_cols : int option; (** Column count for table HDUs. *)\n}\n(** The type for HDU summary information. *)\n\ntype header_card = {\n key : string; (** Keyword name (up to 8 characters). *)\n value : string; (** Parsed value string. *)\n comment : string; (** Inline comment, if any. *)\n}\n(** The type for FITS header cards. *)\n\nval info : string -> hdu_info list\n(** [info path] is the summary information for every HDU in the FITS file at\n [path].\n\n Raises [Failure] if [path] cannot be read or is not a valid FITS file. *)\n\nval header : ?hdu:int -> string -> header_card list\n(** [header path] is the header cards for HDU [hdu] in the FITS file at [path],\n including COMMENT and HISTORY cards.\n\n [hdu] defaults to [0] (primary HDU).\n\n Raises [Failure] if [hdu] is out of range. *)\n\n(** {1:reading Reading} *)\n\nval read_table : ?hdu:int -> string -> Talon.t\n(** [read_table path] reads a BINTABLE extension into a dataframe.\n\n [hdu] defaults to [1] (first extension).\n\n Supported TFORM types: [E] (float32), [D] (float64), [J] (int32), [K]\n (int64), [I] (int16), [B] (uint8), [L] (logical), [A] (string). Vector\n columns (repeat > 1) are not supported except for strings. TSCAL and TZERO\n are applied when present.\n\n Raises [Failure] if the HDU is not a BINTABLE, [hdu] is out of range, or a\n column has an unsupported TFORM type. *)\n\nval read_image : ?hdu:int -> string -> Nx_io.packed\n(** [read_image path] reads an image HDU into a packed {!Nx.t} tensor.\n\n [hdu] defaults to [0] (primary HDU).\n\n Supported BITPIX values: [8], [16], [32], [64], [-32], [-64].\n\n When BSCALE or BZERO header cards are present with non-trivial values\n (BSCALE != 1.0 or BZERO != 0.0), the physical values [BZERO + BSCALE * raw]\n are computed and the result is returned as float64 regardless of the\n original BITPIX. When neither card is present or both have default values,\n the raw data type is preserved.\n\n Raises [Failure] if the HDU is not an image, [hdu] is out of range, or\n BITPIX is unsupported. *)\n\n(** {1:writing Writing} *)\n\nval write_table : ?overwrite:bool -> string -> Talon.t -> unit\n(** [write_table path df] writes [df] as a single BINTABLE extension preceded by\n an empty primary HDU.\n\n [overwrite] defaults to [true].\n\n Raises [Failure] if [overwrite] is [false] and [path] already exists. *)\n\nval write_image : ?overwrite:bool -> string -> ('a, 'b) Nx.t -> unit\n(** [write_image path tensor] writes [tensor] as a primary image HDU.\n\n [overwrite] defaults to [true].\n\n Supported dtypes: uint8, int16, int32, int64, float32, float64.\n\n Raises [Failure] if [overwrite] is [false] and [path] already exists, or if\n the dtype is unsupported. *)\n"
},
{
"path": "dev/umbra/lib/galactocentric.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet pi = Float.pi\nlet f64 = Nx.float64\nlet galcen_distance_default = Unit.Length.of_kpc (Nx.scalar f64 8.122)\nlet z_sun_default = Unit.Length.of_kpc (Nx.scalar f64 0.0208)\n\ntype t = { x : Nx.float64_t; y : Nx.float64_t; z : Nx.float64_t }\n\nlet x t = Unit.Length.of_kpc t.x\nlet y t = Unit.Length.of_kpc t.y\nlet z t = Unit.Length.of_kpc t.z\n\n(* Convert via Galactic coordinates. In Galactic (l,b) the GC is at l=0, b=0, so\n heliocentric Galactic Cartesian is: x_h = d cos(b) cos(l) toward GC y_h = d\n cos(b) sin(l) toward rotation z_h = d sin(b) toward NGP\n\n Galactocentric = heliocentric shifted by Sun's position: x_gc = x_h -\n galcen_distance y_gc = y_h z_gc = z_h + z_sun *)\n\nlet of_coord ?(galcen_distance = galcen_distance_default)\n ?(z_sun = z_sun_default) ~distance c =\n let galcen_distance = Nx.item [] (Unit.Length.in_kpc galcen_distance) in\n let z_sun = Nx.item [] (Unit.Length.in_kpc z_sun) in\n let gal = Coord.galactic c in\n let l_rad = Unit.Angle.to_tensor (Coord.lon gal) in\n let b_rad = Unit.Angle.to_tensor (Coord.lat gal) in\n let d_kpc = Unit.Length.in_kpc distance in\n let n = Nx.numel l_rad in\n let x_out = Nx.zeros Nx.float64 [| n |] in\n let y_out = Nx.zeros Nx.float64 [| n |] in\n let z_out = Nx.zeros Nx.float64 [| n |] in\n for i = 0 to n - 1 do\n let l = Nx.item [ i ] l_rad in\n let b = Nx.item [ i ] b_rad in\n let d = Nx.item [ i ] d_kpc in\n let cb = Float.cos b in\n let xh = d *. cb *. Float.cos l in\n let yh = d *. cb *. Float.sin l in\n let zh = d *. Float.sin b in\n Nx.set_item [ i ] (xh -. galcen_distance) x_out;\n Nx.set_item [ i ] yh y_out;\n Nx.set_item [ i ] (zh +. z_sun) z_out\n done;\n { x = x_out; y = y_out; z = z_out }\n\nlet to_coord ?(galcen_distance = galcen_distance_default)\n ?(z_sun = z_sun_default) t =\n let galcen_distance = Nx.item [] (Unit.Length.in_kpc galcen_distance) in\n let z_sun = Nx.item [] (Unit.Length.in_kpc z_sun) in\n let n = Nx.numel t.x in\n let l_out = Nx.zeros Nx.float64 [| n |] in\n let b_out = Nx.zeros Nx.float64 [| n |] in\n let d_out = Nx.zeros Nx.float64 [| n |] in\n for i = 0 to n - 1 do\n let xg = Nx.item [ i ] t.x in\n let yg = Nx.item [ i ] t.y in\n let zg = Nx.item [ i ] t.z in\n let xh = xg +. galcen_distance in\n let yh = yg in\n let zh = zg -. z_sun in\n let d = Float.sqrt ((xh *. xh) +. (yh *. yh) +. (zh *. zh)) in\n let b = Float.asin (Float.max ~-.1.0 (Float.min 1.0 (zh /. d))) in\n let l = Float.atan2 yh xh in\n let l = if l < 0.0 then l +. (2.0 *. pi) else l in\n Nx.set_item [ i ] l l_out;\n Nx.set_item [ i ] b b_out;\n Nx.set_item [ i ] d d_out\n done;\n let coord =\n Coord.of_galactic\n ~l:(Unit.Angle.of_tensor l_out)\n ~b:(Unit.Angle.of_tensor b_out)\n in\n let distance = Unit.Length.of_kpc d_out in\n (coord, distance)\n"
},
{
"path": "dev/umbra/lib/galactocentric.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Galactocentric Cartesian coordinates.\n\n Converts celestial positions with distances to a right-handed Cartesian\n frame centered on the Galactic center. The x-axis points from the Sun toward\n the Galactic center (l=0, b=0), y in the direction of Galactic rotation, z\n toward the North Galactic Pole.\n\n Coordinates go through the Galactic frame (ICRS {e ->} Galactic {e ->}\n heliocentric Cartesian {e ->} Galactocentric). The Galactic center position\n is defined by the IAU Galactic coordinate system (l=0, b=0).\n\n Default parameters follow\n {{:https://ui.adsabs.harvard.edu/abs/2018A%26A...615L..15G}GRAVITY\n Collaboration (2018)} for the Galactic center distance.\n\n {[\n let star =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 266.0 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| -29.0 |]))\n in\n let gc =\n Galactocentric.of_coord\n ~distance:(Unit.Length.of_kpc (Nx.create f64 [| 1 |] [| 8.0 |]))\n star\n in\n let x_kpc = Nx.item [ 0 ] (Unit.Length.in_kpc (Galactocentric.x gc))\n ]} *)\n\n(** {1:coords Coordinates} *)\n\ntype t\n(** The type for Galactocentric Cartesian positions. *)\n\nval x : t -> Unit.length Unit.t\n(** [x t] is the x coordinate (toward the Galactic center). *)\n\nval y : t -> Unit.length Unit.t\n(** [y t] is the y coordinate (direction of Galactic rotation). *)\n\nval z : t -> Unit.length Unit.t\n(** [z t] is the z coordinate (toward the North Galactic Pole). *)\n\n(** {1:converting Converting} *)\n\nval of_coord :\n ?galcen_distance:Unit.length Unit.t ->\n ?z_sun:Unit.length Unit.t ->\n distance:Unit.length Unit.t ->\n Coord.t ->\n t\n(** [of_coord ~distance c] converts celestial coordinates [c] with [distance] to\n Galactocentric Cartesian. Not differentiable (scalar-level trigonometry).\n\n [galcen_distance] is the Sun-GC distance (defaults to 8.122 kpc, GRAVITY\n Collaboration 2018). [z_sun] is the Sun's height above the Galactic midplane\n (defaults to 0.0208 kpc). *)\n\nval to_coord :\n ?galcen_distance:Unit.length Unit.t ->\n ?z_sun:Unit.length Unit.t ->\n t ->\n Coord.t * Unit.length Unit.t\n(** [to_coord t] converts Galactocentric Cartesian coordinates [t] back to ICRS\n celestial coordinates and a distance. Not differentiable (scalar-level\n trigonometry).\n\n [galcen_distance] defaults to 8.122 kpc. [z_sun] defaults to 0.0208 kpc. *)\n"
},
{
"path": "dev/umbra/lib/kdtree.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype node =\n | Leaf\n | Node of {\n idx : int;\n x : float;\n y : float;\n z : float;\n split : int;\n left : node;\n right : node;\n }\n\ntype t = { root : node; size : int }\n\nlet coord split x y z = match split with 0 -> x | 1 -> y | _ -> z\n\nlet build xs ys zs =\n let n = Array.length xs in\n if n <> Array.length ys || n <> Array.length zs then\n invalid_arg \"Kdtree.build: arrays must have the same length\";\n let indices = Array.init n Fun.id in\n let rec build_rec start len depth =\n if len = 0 then Leaf\n else if len = 1 then\n let i = indices.(start) in\n Node\n {\n idx = i;\n x = xs.(i);\n y = ys.(i);\n z = zs.(i);\n split = depth mod 3;\n left = Leaf;\n right = Leaf;\n }\n else begin\n let split = depth mod 3 in\n let sub = Array.sub indices start len in\n Array.sort\n (fun a b ->\n Float.compare\n (coord split xs.(a) ys.(a) zs.(a))\n (coord split xs.(b) ys.(b) zs.(b)))\n sub;\n Array.blit sub 0 indices start len;\n let mid = len / 2 in\n let mi = indices.(start + mid) in\n let left = build_rec start mid (depth + 1) in\n let right = build_rec (start + mid + 1) (len - mid - 1) (depth + 1) in\n Node\n { idx = mi; x = xs.(mi); y = ys.(mi); z = zs.(mi); split; left; right }\n end\n in\n { root = build_rec 0 n 0; size = n }\n\nlet sq_dist px py pz qx qy qz =\n let dx = px -. qx and dy = py -. qy and dz = pz -. qz in\n (dx *. dx) +. (dy *. dy) +. (dz *. dz)\n\nlet nearest tree qx qy qz =\n if tree.size = 0 then invalid_arg \"Kdtree.nearest: empty tree\";\n let best_idx = ref 0 in\n let best_dist = ref Float.infinity in\n let rec search node =\n match node with\n | Leaf -> ()\n | Node { idx; x; y; z; split; left; right } ->\n let d = sq_dist x y z qx qy qz in\n if d < !best_dist then begin\n best_dist := d;\n best_idx := idx\n end;\n let q_split = coord split qx qy qz in\n let p_split = coord split x y z in\n let diff = q_split -. p_split in\n let near, far = if diff < 0.0 then (left, right) else (right, left) in\n search near;\n if diff *. diff < !best_dist then search far\n in\n search tree.root;\n (!best_idx, !best_dist)\n\nlet within tree qx qy qz max_dist_sq =\n let results = ref [] in\n let rec search node =\n match node with\n | Leaf -> ()\n | Node { idx; x; y; z; split; left; right } ->\n let d = sq_dist x y z qx qy qz in\n if d <= max_dist_sq then results := (idx, d) :: !results;\n let q_split = coord split qx qy qz in\n let p_split = coord split x y z in\n let diff = q_split -. p_split in\n let near, far = if diff < 0.0 then (left, right) else (right, left) in\n search near;\n if diff *. diff <= max_dist_sq then search far\n in\n search tree.root;\n !results\n"
},
{
"path": "dev/umbra/lib/kdtree.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** 3D kd-tree for nearest-neighbor queries.\n\n {b Note.} Private module. *)\n\ntype t\n(** The type for a 3D kd-tree. *)\n\nval build : float array -> float array -> float array -> t\n(** [build xs ys zs] is a kd-tree over the points [(xs.(i), ys.(i), zs.(i))].\n The three arrays must have equal length.\n\n Raises [Invalid_argument] if the arrays differ in length. *)\n\nval nearest : t -> float -> float -> float -> int * float\n(** [nearest tree qx qy qz] is [(i, d2)] where [i] is the index of the nearest\n point to [(qx, qy, qz)] and [d2] is the squared Euclidean distance.\n\n Raises [Invalid_argument] if the tree is empty. *)\n\nval within : t -> float -> float -> float -> float -> (int * float) list\n(** [within tree qx qy qz max_d2] is the list of [(i, d2)] pairs for all points\n within squared Euclidean distance [max_d2] of [(qx, qy, qz)]. *)\n"
},
{
"path": "dev/umbra/lib/photometry.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet f64 = Nx.float64\n\n(* Speed of light for f_lambda to f_nu conversion *)\nlet _c = 299_792_458.0\n\n(* AB magnitude zero-point: 3631 Jy = 3631e-26 W/m²/Hz *)\nlet _ab_zp = 3631.0e-26\n\n(* Wavelength stored internally in metres (SI base unit) *)\ntype bandpass = { wavelength : Nx.float64_t; throughput : Nx.float64_t }\ntype detector = Energy | Photon\n\nlet bandpass ~wavelength ~throughput =\n let wavelength = Unit.Length.to_tensor wavelength in\n if Nx.ndim wavelength <> 1 then\n invalid_arg \"Photometry.bandpass: wavelength must be a 1-D tensor\";\n if Nx.ndim throughput <> 1 then\n invalid_arg \"Photometry.bandpass: throughput must be a 1-D tensor\";\n if Nx.numel wavelength <> Nx.numel throughput then\n invalid_arg\n \"Photometry.bandpass: wavelength and throughput must have the same length\";\n { wavelength; throughput }\n\nlet tophat ~lo ~hi ~n =\n let lo_m = Nx.item [] (Unit.Length.to_tensor lo) in\n let hi_m = Nx.item [] (Unit.Length.to_tensor hi) in\n let wavelength = Nx.linspace f64 lo_m hi_m n in\n let throughput = Nx.ones f64 [| n |] in\n { wavelength; throughput }\n\nlet wavelength bp = Unit.Length.of_tensor bp.wavelength\nlet throughput bp = bp.throughput\n\n(* Differentiable trapezoidal integration along the last axis of y. x is always\n 1-D (the wavelength grid). When y has leading batch dimensions the result\n preserves them. All Nx ops — fully differentiable through Rune. *)\nlet trapz y x =\n let m = Nx.numel x in\n let x0 = Nx.slice [ R (0, m - 1) ] x in\n let x1 = Nx.slice [ R (1, m) ] x in\n let dx = Nx.sub x1 x0 in\n let y_shape = Nx.shape y in\n let ndim = Array.length y_shape in\n if ndim <= 1 then begin\n let y0 = Nx.slice [ R (0, m - 1) ] y in\n let y1 = Nx.slice [ R (1, m) ] y in\n let y_avg = Nx.div_s (Nx.add y0 y1) 2.0 in\n Nx.sum (Nx.mul y_avg dx)\n end\n else begin\n let y2d = Nx.reshape [| -1; m |] y in\n let y0 = Nx.slice [ A; R (0, m - 1) ] y2d in\n let y1 = Nx.slice [ A; R (1, m) ] y2d in\n let y_avg = Nx.div_s (Nx.add y0 y1) 2.0 in\n let result = Nx.sum ~axes:[ 1 ] (Nx.mul y_avg dx) in\n let batch_shape = Array.sub y_shape 0 (ndim - 1) in\n Nx.reshape batch_shape result\n end\n\nlet pivot_wavelength bp =\n let lam = bp.wavelength in\n let t = bp.throughput in\n (* lambda_p = sqrt(integral T lambda d lambda / integral T/lambda d lambda) *)\n let num = trapz (Nx.mul t lam) lam in\n let den = trapz (Nx.div t lam) lam in\n Unit.Length.of_tensor (Nx.sqrt (Nx.div num den))\n\n(* Detector weight: 1 for energy-counting, lambda for photon-counting *)\nlet detector_weight detector lam throughput =\n match detector with Energy -> throughput | Photon -> Nx.mul throughput lam\n\n(* ST magnitude zero-point: -2.5 log10(f_lambda / 3.63e-9 erg/s/cm²/Å) In SI:\n 3.63e-9 erg/s/cm²/Å = 3.63e-9 * 1e-7 * 1e4 * 1e10 W/m²/m = 3.63e-2 W/m²/m *)\nlet _st_zp = 3.63e-2\n\nlet align_spectrum bp spectrum =\n let lam_bp = bp.wavelength in\n let lam_sp = Unit.Length.to_tensor (Spectrum.wavelength spectrum) in\n let same =\n Nx.numel lam_bp = Nx.numel lam_sp\n && Nx.item [] (Nx.max (Nx.abs (Nx.sub lam_bp lam_sp))) = 0.0\n in\n if same then spectrum\n else Spectrum.resample ~wavelength:(Unit.Length.of_tensor lam_bp) spectrum\n\nlet flux_density ?(detector = Energy) bp spectrum =\n let spectrum = align_spectrum bp spectrum in\n let lam = bp.wavelength in\n let f = Spectrum.values spectrum in\n let w = detector_weight detector lam bp.throughput in\n Nx.div (trapz (Nx.mul f w) lam) (trapz w lam)\n\nlet ab_mag ?(detector = Energy) bp spectrum =\n let spectrum = align_spectrum bp spectrum in\n let lam = bp.wavelength in\n let f_lambda = Spectrum.values spectrum in\n let f_nu = Nx.div (Nx.mul f_lambda (Nx.square lam)) (Nx.scalar f64 _c) in\n let w = detector_weight detector lam bp.throughput in\n let mean_fnu = Nx.div (trapz (Nx.mul f_nu w) lam) (trapz w lam) in\n Nx.mul_s\n (Nx.log (Nx.div mean_fnu (Nx.scalar f64 _ab_zp)))\n (-2.5 /. Float.log 10.0)\n\nlet st_mag ?(detector = Energy) bp spectrum =\n let spectrum = align_spectrum bp spectrum in\n let lam = bp.wavelength in\n let f_lambda = Spectrum.values spectrum in\n let w = detector_weight detector lam bp.throughput in\n let mean_flam = Nx.div (trapz (Nx.mul f_lambda w) lam) (trapz w lam) in\n Nx.mul_s\n (Nx.log (Nx.div mean_flam (Nx.scalar f64 _st_zp)))\n (-2.5 /. Float.log 10.0)\n\nlet _vega_spectrum =\n let n = Array.length Vega_data.wave in\n let w = Nx.create f64 [| n |] Vega_data.wave in\n let w = Nx.mul_s w 1e-10 in\n let f = Nx.create f64 [| n |] Vega_data.flux in\n Spectrum.create ~wavelength:(Unit.Length.of_tensor w) ~values:f\n |> Spectrum.as_flux_density\n\nlet vega_mag ?(detector = Energy) bp spectrum =\n let f_src = flux_density ~detector bp spectrum in\n let f_vega = flux_density ~detector bp _vega_spectrum in\n Nx.mul_s (Nx.log (Nx.div f_src f_vega)) (-2.5 /. Float.log 10.0)\n\nlet color ?detector bp1 bp2 spectrum =\n Nx.sub (ab_mag ?detector bp1 spectrum) (ab_mag ?detector bp2 spectrum)\n\nlet effective_wavelength ?(detector = Energy) bp spectrum =\n let spectrum = align_spectrum bp spectrum in\n let lam = bp.wavelength in\n let f = Spectrum.values spectrum in\n let w = detector_weight detector lam bp.throughput in\n let fw = Nx.mul f w in\n let num = trapz (Nx.mul fw (Nx.square lam)) lam in\n let den = trapz (Nx.mul fw lam) lam in\n Unit.Length.of_tensor (Nx.div num den)\n"
},
{
"path": "dev/umbra/lib/photometry.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Synthetic photometry.\n\n Computes broadband fluxes and magnitudes by integrating spectra through\n filter bandpasses using trapezoidal quadrature.\n\n {[\n let bp = Photometry.tophat\n ~lo:(Unit.Length.nm 400.0)\n ~hi:(Unit.Length.nm 500.0) ~n:100 in\n let mag = Photometry.ab_mag bp sed\n ]}\n\n All photometry functions accept batched spectra (values with leading batch\n dimensions). When a spectrum has shape [[batch; n_lambda]], the result has\n shape [[batch]]. *)\n\n(** {1:types Types} *)\n\ntype bandpass\n(** The type for filter transmission curves. *)\n\ntype detector =\n | Energy\n | Photon\n (** The detector convention.\n\n - {!Energy}: counts incident energy (default). The bandpass-weighted\n mean is [ = integral f_nu T d lambda / integral T d lambda].\n - {!Photon}: counts photons. Weights both numerator and denominator by\n [lambda]:\n [ = integral f_nu T lambda d lambda / integral T lambda d\n lambda]. *)\n\n(** {1:constructors Constructors} *)\n\nval bandpass :\n wavelength:Unit.length Unit.t -> throughput:Nx.float64_t -> bandpass\n(** [bandpass ~wavelength ~throughput] is a filter from 1-D arrays. [throughput]\n is dimensionless (typically in \\[0, 1\\]).\n\n Raises [Invalid_argument] if tensors are not 1-D or have different lengths.\n*)\n\nval tophat : lo:Unit.length Unit.t -> hi:Unit.length Unit.t -> n:int -> bandpass\n(** [tophat ~lo ~hi ~n] is a rectangular bandpass from [lo] to [hi] with [n]\n wavelength points and unit throughput. *)\n\n(** {1:accessors Accessors} *)\n\nval wavelength : bandpass -> Unit.length Unit.t\n(** [wavelength bp] is the wavelength grid. *)\n\nval throughput : bandpass -> Nx.float64_t\n(** [throughput bp] is the throughput curve. *)\n\nval pivot_wavelength : bandpass -> Unit.length Unit.t\n(** [pivot_wavelength bp] is the pivot wavelength\n {e lambda}{_ p}[ = sqrt(integral T lambda d lambda / integral T/lambda d\n lambda)]. *)\n\n(** {1:photometry Synthetic photometry} *)\n\nval flux_density :\n ?detector:detector ->\n bandpass ->\n Spectrum.flux_density Spectrum.t ->\n Nx.float64_t\n(** [flux_density ?detector bp spectrum] is the bandpass-weighted mean flux\n density [ = integral f T w d lambda / integral T w d lambda] where [w] is\n [1] for {!Energy} and [lambda] for {!Photon}. [detector] defaults to\n {!Energy}.\n\n The spectrum is resampled to the bandpass wavelength grid via linear\n interpolation if they differ. Differentiable through Rune. *)\n\nval ab_mag :\n ?detector:detector ->\n bandpass ->\n Spectrum.flux_density Spectrum.t ->\n Nx.float64_t\n(** [ab_mag ?detector bp spectrum] is the AB magnitude of [spectrum] through\n [bp].\n\n Computes the mean spectral flux density in f{_ nu}:\n [ = integral (f_lambda lambda{^2}/c) T w d lambda / integral T w d\n lambda], where [w] is [1] for {!Energy} and [lambda] for {!Photon}, then\n [m_AB = -2.5 log10( / 3631 Jy)]. [detector] defaults to {!Energy}.\n\n The spectrum is resampled to the bandpass wavelength grid via linear\n interpolation if they differ. Differentiable through Rune. *)\n\nval st_mag :\n ?detector:detector ->\n bandpass ->\n Spectrum.flux_density Spectrum.t ->\n Nx.float64_t\n(** [st_mag ?detector bp spectrum] is the ST magnitude of [spectrum] through\n [bp].\n\n Computes the bandpass-weighted mean f{_ lambda}, then\n [m_ST = -2.5 log10( / 3.63e-9 erg s{^-1} cm{^-2} A{^-1})].\n [detector] defaults to {!Energy}.\n\n The spectrum is resampled to the bandpass wavelength grid via linear\n interpolation if they differ. Differentiable through Rune. *)\n\nval vega_mag :\n ?detector:detector ->\n bandpass ->\n Spectrum.flux_density Spectrum.t ->\n Nx.float64_t\n(** [vega_mag ?detector bp spectrum] is the Vega magnitude of [spectrum] through\n [bp].\n\n Computes [-2.5 log10( / )] where the Vega reference\n spectrum is from CALSPEC alpha_lyr_stis_011.fits (Bohlin 2014). [detector]\n defaults to {!Energy}.\n\n The spectrum is resampled to the bandpass wavelength grid via linear\n interpolation if they differ. Differentiable through Rune. *)\n\nval color :\n ?detector:detector ->\n bandpass ->\n bandpass ->\n Spectrum.flux_density Spectrum.t ->\n Nx.float64_t\n(** [color ?detector bp1 bp2 spectrum] is\n [ab_mag ?detector bp1 spectrum - ab_mag ?detector bp2 spectrum].\n\n Differentiable through Rune. *)\n\nval effective_wavelength :\n ?detector:detector ->\n bandpass ->\n Spectrum.flux_density Spectrum.t ->\n Unit.length Unit.t\n(** [effective_wavelength ?detector bp spectrum] is the source-dependent\n effective wavelength\n {e lambda}{_ eff}[ = integral f T w lambda{^2} d lambda / integral f T w\n lambda d lambda].\n\n Unlike {!pivot_wavelength}, this depends on the source spectrum. The\n spectrum is resampled if grids differ. Differentiable through Rune. *)\n"
},
{
"path": "dev/umbra/lib/spectrum.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet f64 = Nx.float64\n\n(* Physical constants (SI) *)\nlet _h = 6.626_070_15e-34\nlet _c = 299_792_458.0\nlet _k_b = 1.380_649e-23\nlet _two_hc2 = 2.0 *. _h *. _c *. _c\nlet _hc_over_k = _h *. _c /. _k_b\n\n(* Spectral kinds — phantom, no runtime representation *)\ntype flux_density\ntype radiance\ntype sampled\n\n(* Wavelength stored internally in metres (SI base unit) *)\ntype 'a t = { wavelength : Nx.float64_t; values : Nx.float64_t }\n\nlet validate_increasing name wl =\n let n = Nx.numel wl in\n if n > 1 then\n for i = 1 to n - 1 do\n if Nx.item [ i ] wl <= Nx.item [ i - 1 ] wl then\n invalid_arg (name ^ \": wavelength must be strictly increasing\")\n done\n\nlet create ~wavelength ~values =\n let wavelength = Unit.Length.to_tensor wavelength in\n if Nx.ndim wavelength <> 1 then\n invalid_arg \"Spectrum.create: wavelength must be a 1-D tensor\";\n let v_shape = Nx.shape values in\n let v_ndim = Array.length v_shape in\n if v_ndim = 0 then invalid_arg \"Spectrum.create: values must be at least 1-D\";\n if v_shape.(v_ndim - 1) <> Nx.numel wavelength then\n invalid_arg\n \"Spectrum.create: last dimension of values must match wavelength length\";\n validate_increasing \"Spectrum.create\" wavelength;\n { wavelength; values }\n\nlet wavelength t = Unit.Length.of_tensor t.wavelength\nlet values t = t.values\nlet as_flux_density t = { wavelength = t.wavelength; values = t.values }\nlet as_sampled t = { wavelength = t.wavelength; values = t.values }\n\nlet blackbody ~temperature ~wavelength =\n let wavelength = Unit.Length.to_tensor wavelength in\n let temp = Unit.Temperature.to_tensor temperature in\n let two_hc2 = Nx.scalar f64 _two_hc2 in\n let hc_k = Nx.scalar f64 _hc_over_k in\n let lam5 = Nx.pow_s wavelength 5.0 in\n let exponent = Nx.div hc_k (Nx.mul wavelength temp) in\n let values =\n Nx.div (Nx.div two_hc2 lam5) (Nx.sub (Nx.exp exponent) (Nx.scalar f64 1.0))\n in\n { wavelength; values }\n\nlet power_law ~amplitude ~index ~pivot ~wavelength =\n let wavelength = Unit.Length.to_tensor wavelength in\n let pivot = Unit.Length.to_tensor pivot in\n let ratio = Nx.div wavelength pivot in\n let values = Nx.mul amplitude (Nx.pow ratio index) in\n { wavelength; values }\n\nlet redshift ~z t =\n let one_plus_z = Nx.add_s z 1.0 in\n {\n wavelength = Nx.mul t.wavelength one_plus_z;\n values = Nx.div t.values one_plus_z;\n }\n\nlet scale factor t = { t with values = Nx.mul factor t.values }\n\nlet mul a b =\n if Nx.numel a.wavelength <> Nx.numel b.wavelength then\n invalid_arg \"Spectrum.mul: spectra must have the same wavelength grid\";\n let max_diff =\n Nx.item [] (Nx.max (Nx.abs (Nx.sub a.wavelength b.wavelength)))\n in\n if max_diff > 0.0 then\n invalid_arg \"Spectrum.mul: spectra must have the same wavelength grid\";\n { wavelength = a.wavelength; values = Nx.mul a.values b.values }\n\nlet div a b =\n if Nx.numel a.wavelength <> Nx.numel b.wavelength then\n invalid_arg \"Spectrum.div: spectra must have the same wavelength grid\";\n let max_diff =\n Nx.item [] (Nx.max (Nx.abs (Nx.sub a.wavelength b.wavelength)))\n in\n if max_diff > 0.0 then\n invalid_arg \"Spectrum.div: spectra must have the same wavelength grid\";\n { wavelength = a.wavelength; values = Nx.div a.values b.values }\n\nlet add a b =\n if Nx.numel a.wavelength <> Nx.numel b.wavelength then\n invalid_arg \"Spectrum.add: spectra must have the same wavelength grid\";\n let max_diff =\n Nx.item [] (Nx.max (Nx.abs (Nx.sub a.wavelength b.wavelength)))\n in\n if max_diff > 0.0 then\n invalid_arg \"Spectrum.add: spectra must have the same wavelength grid\";\n { wavelength = a.wavelength; values = Nx.add a.values b.values }\n\nlet resample ~wavelength t =\n let new_wave = Unit.Length.to_tensor wavelength in\n if Nx.ndim new_wave <> 1 then\n invalid_arg \"Spectrum.resample: wavelength must be a 1-D tensor\";\n validate_increasing \"Spectrum.resample\" new_wave;\n let old_wave = t.wavelength in\n let old_values = t.values in\n let n_old = Nx.numel old_wave and n_new = Nx.numel new_wave in\n (* Find lower bracket index for each target wavelength (non-differentiable) *)\n let lo_arr =\n Array.init n_new (fun j ->\n let x = Nx.item [ j ] new_wave in\n let lo = ref 0 and hi = ref (n_old - 1) in\n while !hi - !lo > 1 do\n let mid = (!lo + !hi) / 2 in\n if Nx.item [ mid ] old_wave <= x then lo := mid else hi := mid\n done;\n !lo)\n in\n let hi_arr =\n Array.init n_new (fun j -> Int32.of_int (min (lo_arr.(j) + 1) (n_old - 1)))\n in\n let lo_arr = Array.map Int32.of_int lo_arr in\n let lo_t = Nx.create Nx.int32 [| n_new |] lo_arr in\n let hi_t = Nx.create Nx.int32 [| n_new |] hi_arr in\n (* Gather source wavelengths and values at bracket endpoints. Nx.take uses\n B.gather, which Rune differentiates through. *)\n let x0 = Nx.take lo_t old_wave in\n let x1 = Nx.take hi_t old_wave in\n let y0 = Nx.take ~axis:(-1) lo_t old_values in\n let y1 = Nx.take ~axis:(-1) hi_t old_values in\n (* Linear interpolation — differentiable through Rune *)\n let dx = Nx.clamp ~min:1e-30 (Nx.sub x1 x0) in\n let frac = Nx.div (Nx.sub new_wave x0) dx in\n let values = Nx.add y0 (Nx.mul frac (Nx.sub y1 y0)) in\n { wavelength = new_wave; values }\n\nlet gaussian ~amplitude ~center ~stddev ~wavelength =\n let wavelength = Unit.Length.to_tensor wavelength in\n let center = Unit.Length.to_tensor center in\n let stddev = Unit.Length.to_tensor stddev in\n let x = Nx.sub wavelength center in\n let z = Nx.div x stddev in\n let values = Nx.mul amplitude (Nx.exp (Nx.mul_s (Nx.mul z z) (-0.5))) in\n { wavelength; values }\n\nlet lorentzian ~amplitude ~center ~fwhm ~wavelength =\n let wavelength = Unit.Length.to_tensor wavelength in\n let center = Unit.Length.to_tensor center in\n let half_gamma = Nx.div_s (Unit.Length.to_tensor fwhm) 2.0 in\n let x = Nx.sub wavelength center in\n let hg2 = Nx.mul half_gamma half_gamma in\n let values = Nx.mul amplitude (Nx.div hg2 (Nx.add (Nx.mul x x) hg2)) in\n { wavelength; values }\n\nlet voigt ~amplitude ~center ~sigma ~gamma ~wavelength =\n let wavelength = Unit.Length.to_tensor wavelength in\n let center = Unit.Length.to_tensor center in\n let sigma = Unit.Length.to_tensor sigma in\n let gamma = Unit.Length.to_tensor gamma in\n (* Pseudo-Voigt mixing via Thompson, Cox & Hastings (1987). *)\n let sqrt_2ln2 = Float.sqrt (2.0 *. Float.log 2.0) in\n let fg = Nx.mul_s sigma (2.0 *. sqrt_2ln2) in\n let fl = Nx.mul_s gamma 2.0 in\n let fg2 = Nx.mul fg fg in\n let fg3 = Nx.mul fg2 fg in\n let fg4 = Nx.mul fg3 fg in\n let fg5 = Nx.mul fg4 fg in\n let fl2 = Nx.mul fl fl in\n let fl3 = Nx.mul fl2 fl in\n let fl4 = Nx.mul fl3 fl in\n let fl5 = Nx.mul fl4 fl in\n let f =\n Nx.pow_s\n (Nx.add fg5\n (Nx.add\n (Nx.mul_s (Nx.mul fg4 fl) 2.69269)\n (Nx.add\n (Nx.mul_s (Nx.mul fg3 fl2) 2.42843)\n (Nx.add\n (Nx.mul_s (Nx.mul fg2 fl3) 4.47163)\n (Nx.add (Nx.mul_s (Nx.mul fg fl4) 0.07842) fl5)))))\n 0.2\n in\n let ratio = Nx.div fl f in\n let ratio2 = Nx.mul ratio ratio in\n let ratio3 = Nx.mul ratio2 ratio in\n let eta =\n Nx.add (Nx.mul_s ratio 1.36603)\n (Nx.add (Nx.mul_s ratio2 (-0.47719)) (Nx.mul_s ratio3 0.11116))\n in\n (* Gaussian component (unit height at center) *)\n let x = Nx.sub wavelength center in\n let sig_eff = Nx.div_s f (2.0 *. sqrt_2ln2) in\n let z_g = Nx.div x sig_eff in\n let gauss = Nx.exp (Nx.mul_s (Nx.mul z_g z_g) (-0.5)) in\n (* Lorentzian component (unit height at center) *)\n let hf = Nx.div_s f 2.0 in\n let hf2 = Nx.mul hf hf in\n let lorentz = Nx.div hf2 (Nx.add (Nx.mul x x) hf2) in\n let values =\n Nx.mul amplitude\n (Nx.add (Nx.mul eta lorentz)\n (Nx.mul (Nx.sub (Nx.scalar f64 1.0) eta) gauss))\n in\n { wavelength; values }\n"
},
{
"path": "dev/umbra/lib/spectrum.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Sampled spectral quantities on a wavelength grid.\n\n A {!'a t} pairs a wavelength grid with spectral values parameterised by a\n phantom {e kind} that tracks the physical meaning of the values:\n\n - {!flux_density}: spectral flux density f{_ lambda} (W m{^ -2} m{^ -1}).\n - {!radiance}: spectral radiance B{_ lambda} (W m{^ -2} m{^ -1} sr{^ -1}).\n - {!sampled}: arbitrary values with no physical assumption.\n\n Operations that depend on the physical interpretation of the values (e.g.,\n {!redshift}, {!val-Photometry.ab_mag}) require a specific kind, preventing\n accidental misuse at compile time. Use {!as_flux_density} to explicitly\n reinterpret values when the physical meaning is known to the caller.\n\n {[\n let wave = Unit.Length.of_m (Nx.linspace Nx.float64 1e-7 1e-5 1000) in\n let sed =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar Nx.float64 5800.0))\n ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let reddened =\n Extinction.apply (Extinction.ccm89 ~rv) ~av sed\n ]}\n\n {2:batch Batched spectra}\n\n Values may have leading batch dimensions: a spectrum with wavelength\n [[n_lambda]] and values [[batch; n_lambda]] represents [batch] spectra\n sharing a wavelength grid. All operations ({!resample}, {!scale}, {!add},\n {!val-Photometry.ab_mag}, {!val-Extinction.apply}, etc.) broadcast over\n leading dimensions via Nx:\n\n {[\n let values = Nx.stack (List.map Spectrum.values templates) in\n let batch =\n Spectrum.create ~wavelength ~values |> Spectrum.as_flux_density\n in\n let mags = Photometry.ab_mag bp batch (* shape [batch] *)\n ]}\n\n {b Note.} {!redshift} with a per-spectrum [z] does not broadcast — it\n changes the wavelength grid, breaking the shared-grid invariant. Use\n [List.map] or [Rune.vmap] for per-spectrum redshifts. *)\n\n(** {1:kinds Spectral kinds} *)\n\ntype flux_density\n(** Phantom type for spectral flux density f{_ lambda} (W m{^ -2} m{^ -1}). *)\n\ntype radiance\n(** Phantom type for spectral radiance B{_ lambda} (W m{^ -2} m{^ -1} sr{^ -1}).\n*)\n\ntype sampled\n(** Phantom type for arbitrary sampled spectral values. *)\n\n(** {1:types Types} *)\n\ntype 'a t\n(** The type for spectra parameterised by spectral kind ['a]. *)\n\n(** {1:constructors Constructors} *)\n\nval create : wavelength:Unit.length Unit.t -> values:Nx.float64_t -> sampled t\n(** [create ~wavelength ~values] is a tabulated spectrum. [wavelength] must be\n 1-D. [values] must be at least 1-D with its last dimension matching\n [wavelength]; leading dimensions are preserved as batch dimensions.\n\n Raises [Invalid_argument] if [wavelength] is not 1-D, the last dimension of\n [values] does not match, or [wavelength] is not strictly increasing. *)\n\n(** {1:accessors Accessors} *)\n\nval wavelength : 'a t -> Unit.length Unit.t\n(** [wavelength s] is the wavelength grid. *)\n\nval values : 'a t -> Nx.float64_t\n(** [values s] is the spectral values. *)\n\n(** {1:casts Kind casts} *)\n\nval as_flux_density : _ t -> flux_density t\n(** [as_flux_density s] reinterprets [s] as spectral flux density. The caller is\n responsible for ensuring the values represent f{_ lambda}. Use this when\n working with external data or when only relative values matter (e.g.,\n fitting colours from a blackbody model). *)\n\nval as_sampled : _ t -> sampled t\n(** [as_sampled s] forgets the spectral kind. *)\n\n(** {1:models Parametric models} *)\n\nval blackbody :\n temperature:Unit.temperature Unit.t ->\n wavelength:Unit.length Unit.t ->\n radiance t\n(** [blackbody ~temperature ~wavelength] is the Planck spectral radiance\n B{_ lambda}(T) in W m{^ -2} m{^ -1} sr{^ -1} at the given wavelengths. This\n is a per-steradian quantity; multiply by a solid angle to obtain spectral\n irradiance. Differentiable through Rune. *)\n\nval power_law :\n amplitude:Nx.float64_t ->\n index:Nx.float64_t ->\n pivot:Unit.length Unit.t ->\n wavelength:Unit.length Unit.t ->\n sampled t\n(** [power_law ~amplitude ~index ~pivot ~wavelength] is the spectrum\n [amplitude * (wavelength / pivot){^index}]. Differentiable through Rune. *)\n\n(** {1:operations Operations} *)\n\nval redshift : z:Nx.float64_t -> flux_density t -> flux_density t\n(** [redshift ~z s] shifts [s] to redshift [z]. Wavelengths are multiplied by\n [(1+z)] and values are divided by [(1+z)].\n\n Restricted to {!flux_density} spectra because the [(1+z){^ -1}] dimming\n factor is specific to spectral flux density. Differentiable through Rune. *)\n\nval scale : Nx.float64_t -> 'a t -> 'a t\n(** [scale factor s] is [s] with values multiplied element-wise by [factor].\n [factor] may be a scalar or a tensor that broadcasts with the values.\n Differentiable through Rune. *)\n\nval mul : 'a t -> sampled t -> 'a t\n(** [mul a b] multiplies values element-wise. [a]'s spectral kind is preserved;\n [b] is treated as a dimensionless modifier (transmission curve, efficiency\n function, etc.). Both must share the same wavelength grid. Differentiable\n through Rune.\n\n Raises [Invalid_argument] if wavelength grids have different lengths. *)\n\nval div : 'a t -> sampled t -> 'a t\n(** [div a b] divides values element-wise. [a]'s spectral kind is preserved; [b]\n is treated as a dimensionless modifier. Both must share the same wavelength\n grid. Differentiable through Rune.\n\n Raises [Invalid_argument] if wavelength grids have different lengths. *)\n\nval add : 'a t -> 'a t -> 'a t\n(** [add a b] is the element-wise sum of two spectra. Both must share the same\n wavelength grid. Differentiable through Rune.\n\n Raises [Invalid_argument] if wavelength grids have different lengths. *)\n\nval resample : wavelength:Unit.length Unit.t -> 'a t -> 'a t\n(** [resample ~wavelength s] resamples [s] onto a new wavelength grid using\n linear interpolation. Leading batch dimensions are preserved. Differentiable\n through Rune with respect to the spectrum values (index computation is not\n differentiable, but the interpolation weights and gather operations are).\n\n Raises [Invalid_argument] if [wavelength] is not 1-D or not strictly\n increasing. *)\n\n(** {1:lines Line profiles} *)\n\nval gaussian :\n amplitude:Nx.float64_t ->\n center:Unit.length Unit.t ->\n stddev:Unit.length Unit.t ->\n wavelength:Unit.length Unit.t ->\n sampled t\n(** [gaussian ~amplitude ~center ~stddev ~wavelength] is the Gaussian profile\n [amplitude * exp(-0.5 * ((lambda - center) / stddev){^2})].\n\n [amplitude], [center], and [stddev] may be scalar tensors; they broadcast\n against [wavelength]. Differentiable through Rune. *)\n\nval lorentzian :\n amplitude:Nx.float64_t ->\n center:Unit.length Unit.t ->\n fwhm:Unit.length Unit.t ->\n wavelength:Unit.length Unit.t ->\n sampled t\n(** [lorentzian ~amplitude ~center ~fwhm ~wavelength] is the Lorentzian profile\n [amplitude * (gamma/2){^2} / ((lambda - center){^2} + (gamma/2){^2})] where\n [gamma = fwhm]. Unit height at [center]. Differentiable through Rune. *)\n\nval voigt :\n amplitude:Nx.float64_t ->\n center:Unit.length Unit.t ->\n sigma:Unit.length Unit.t ->\n gamma:Unit.length Unit.t ->\n wavelength:Unit.length Unit.t ->\n sampled t\n(** [voigt ~amplitude ~center ~sigma ~gamma ~wavelength] is the pseudo-Voigt\n approximation of the Voigt profile (Thompson, Cox & Hastings 1987). [sigma]\n is the Gaussian standard deviation and [gamma] is the Lorentzian half-width\n at half-maximum. Accurate to <1% of the exact Faddeeva-based Voigt.\n Differentiable through Rune. *)\n"
},
{
"path": "dev/umbra/lib/survey.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet f64 = Nx.float64\nlet c_km_s = 299792.458\nlet h0_ref = 100.0\nlet steradian_to_arcmin2 = 11818102.86004228\nlet c1_rho_crit = 0.0134\n\n(* Redshift distributions *)\n\ntype nz = { eval : Nx.float64_t -> Nx.float64_t; zmax : float }\n\nlet simps_float f a b n =\n let h = (b -. a) /. Float.of_int n in\n let sum = ref (f a +. f b) in\n for i = 1 to n - 1 do\n let x = a +. (Float.of_int i *. h) in\n let w = if i mod 2 = 1 then 4.0 else 2.0 in\n sum := !sum +. (w *. f x)\n done;\n !sum *. h /. 3.0\n\nlet smail ?(zmax = 10.0) ~a ~b ~z0 () =\n let raw z_f = (z_f ** a) *. Float.exp (-.((z_f /. z0) ** b)) in\n let norm = simps_float raw 0.0 zmax 256 in\n let eval z =\n let z_f = Nx.item [] z in\n Nx.scalar f64 (raw z_f /. norm)\n in\n { eval; zmax }\n\nlet tabulated ~z ~pz () =\n let n = (Nx.shape z).(0) in\n let zmax = Nx.item [ n - 1 ] z in\n let norm = ref 0.0 in\n for i = 0 to n - 2 do\n let dz = Nx.item [ i + 1 ] z -. Nx.item [ i ] z in\n norm := !norm +. (0.5 *. (Nx.item [ i ] pz +. Nx.item [ i + 1 ] pz) *. dz)\n done;\n let eval zq =\n let zq_f = Nx.item [] zq in\n if zq_f <= Nx.item [ 0 ] z || zq_f >= zmax then Nx.scalar f64 0.0\n else begin\n let idx = ref 0 in\n for i = 0 to n - 2 do\n if Nx.item [ i ] z <= zq_f then idx := i\n done;\n let i = !idx in\n let z0 = Nx.item [ i ] z and z1 = Nx.item [ i + 1 ] z in\n let p0 = Nx.item [ i ] pz and p1 = Nx.item [ i + 1 ] pz in\n let frac = (zq_f -. z0) /. (z1 -. z0) in\n Nx.scalar f64 ((p0 +. (frac *. (p1 -. p0))) /. !norm)\n end\n in\n { eval; zmax }\n\nlet custom_nz ?(zmax = 10.0) eval = { eval; zmax }\nlet eval_nz nz z = nz.eval z\nlet nz_zmax nz = nz.zmax\n\n(* Galaxy bias *)\n\ntype bias = Cosmo.params -> Nx.float64_t -> Nx.float64_t\n\nlet constant_bias b _p _z = Nx.scalar f64 b\n\nlet inverse_growth_bias b0 p z =\n let d = Cosmo.growth_factor ~p z in\n Nx.div (Nx.scalar f64 b0) d\n\n(* Power spectrum backends *)\n\ntype power = Cosmo.params -> Nx.float64_t -> Nx.float64_t -> Nx.float64_t\n\nlet linear p k z = Cosmo.linear_power ~p k z\nlet nonlinear p k z = Cosmo.nonlinear_power ~p k z\n\nlet baryonic_feedback ?(a_bary = 0.0) ?(log10_k_star = 1.0) ?(sigma = 0.55)\n base_power =\n fun p k z ->\n let pk = base_power p k z in\n if a_bary = 0.0 then pk\n else\n let inv_sigma2 = -1.0 /. (sigma *. sigma) in\n let log10_k = Nx.div_s (Nx.log k) (Float.log 10.0) in\n let delta = Nx.sub_s log10_k log10_k_star in\n let gauss = Nx.exp (Nx.mul_s (Nx.mul delta delta) inv_sigma2) in\n Nx.sub pk (Nx.mul_s (Nx.mul gauss pk) a_bary)\n\n(* Tracers *)\n\ntype tracer_kind =\n | Weak_lensing of { ia_bias : bias option; sigma_e : float; m_bias : float }\n | Number_counts of { bias : bias }\n | Custom of {\n kernel :\n p:Cosmo.params -> z:Nx.float64_t -> chi:Nx.float64_t -> Nx.float64_t;\n }\n\ntype tracer = {\n nz : nz option;\n n_gal : float;\n noise : float;\n kind : tracer_kind;\n zmax : float;\n}\n\nlet weak_lensing ?ia_bias ?(sigma_e = 0.26) ?(m_bias = 0.0) ?(n_gal = 1.0) nz =\n let noise = sigma_e *. sigma_e /. (n_gal *. steradian_to_arcmin2) in\n {\n nz = Some nz;\n n_gal;\n noise;\n kind = Weak_lensing { ia_bias; sigma_e; m_bias };\n zmax = nz.zmax;\n }\n\nlet number_counts ~bias ?(n_gal = 1.0) nz =\n let noise = 1.0 /. (n_gal *. steradian_to_arcmin2) in\n { nz = Some nz; n_gal; noise; kind = Number_counts { bias }; zmax = nz.zmax }\n\nlet tracer ?(noise = 0.0) ?(zmax = 3.0) kernel =\n { nz = None; n_gal = 0.0; noise; kind = Custom { kernel }; zmax }\n\n(* Cls result type *)\n\ntype cls = {\n ell : Nx.float64_t;\n tracers : tracer array;\n spectra : Nx.float64_t;\n}\n\n(* Cl index ordering: upper triangle *)\n\nlet pair_index nt i j =\n let a, b = if i <= j then (i, j) else (j, i) in\n (a * ((2 * nt) - a - 1) / 2) + b\n\nlet cl_pairs nt =\n let pairs = ref [] in\n for i = 0 to nt - 1 do\n for j = i to nt - 1 do\n pairs := (i, j) :: !pairs\n done\n done;\n List.rev !pairs\n\n(* Evaluate n(z) for one bin on the z grid. Returns tensor [n_z]. Uses Nx.stack\n so gradients flow through custom_nz eval functions. *)\nlet eval_nz_grid nz z_arr n_z =\n Nx.stack (List.init n_z (fun j -> nz.eval (Nx.scalar f64 z_arr.(j))))\n\n(* Reverse cumulative trapezoidal sum of tensor [n] with spacing dz. result[j] =\n ∫_{x_j}^{x_{n-1}} f(x) dx via trapezoidal rule. *)\nlet rev_cumtrapz f_vec n dz =\n let left = Nx.slice [ R (0, n - 1) ] f_vec in\n let right = Nx.slice [ R (1, n) ] f_vec in\n let mid = Nx.mul_s (Nx.add left right) (0.5 *. dz) in\n let partial = Nx.flip (Nx.cumsum ~axis:0 (Nx.flip mid)) in\n Nx.concatenate [ partial; Nx.zeros f64 [| 1 |] ]\n\n(* Angular power spectra *)\n\nlet angular_cl ?(p = Cosmo.planck18) ?(power = nonlinear) ~ell tracers =\n let tracers_arr = Array.of_list tracers in\n let nt = Array.length tracers_arr in\n let pairs = cl_pairs nt in\n let pairs_arr = Array.of_list pairs in\n let zmax =\n Array.fold_left (fun acc t -> Float.max acc t.zmax) 0.0 tracers_arr\n in\n let n_z = 100 in\n let dz = zmax /. Float.of_int (n_z - 1) in\n let z_arr = Array.init n_z (fun i -> Float.of_int i *. dz) in\n z_arr.(0) <- 1e-6;\n let z_vec = Nx.create f64 [| n_z |] z_arr in\n\n (* Simpson weights: tensor [n_z] *)\n let sw =\n Array.init n_z (fun i ->\n if i = 0 || i = n_z - 1 then 1.0 else if i mod 2 = 1 then 4.0 else 2.0)\n in\n let simpson_w = Nx.mul_s (Nx.create f64 [| n_z |] sw) (dz /. 3.0) in\n\n (* Precompute z-dependent quantities as tensors — differentiable through p.\n comoving_distance and growth_factor use GL quadrature internally and cannot\n accept vector z, so we loop over scalar z values. *)\n let h_t = Nx.div (Cosmo.h0 p) (Nx.scalar f64 h0_ref) in\n let chi_vec =\n Nx.stack\n (List.init n_z (fun j ->\n let z_t = Nx.scalar f64 z_arr.(j) in\n Nx.mul (Unit.Length.in_mpc (Cosmo.comoving_distance ~p z_t)) h_t))\n in\n let chi_safe = Nx.clamp ~min:1e-10 chi_vec in\n let h_vec = Cosmo.hubble ~p z_vec in\n let dchi_dz_vec = Nx.div (Nx.mul_s h_t c_km_s) h_vec in\n let growth_vec =\n Nx.stack\n (List.init n_z (fun j -> Cosmo.growth_factor ~p (Nx.scalar f64 z_arr.(j))))\n in\n let omega_m_t = Cosmo.omega_m p in\n\n (* n(z) values per tracer: tensors [n_z], differentiable through custom_nz *)\n let nz_arrs = Array.make nt (Nx.zeros f64 [| n_z |]) in\n Array.iteri\n (fun idx t ->\n match t.nz with\n | Some nz -> nz_arrs.(idx) <- eval_nz_grid nz z_arr n_z\n | None -> ())\n tracers_arr;\n\n (* Kernel base vectors per tracer: tensor [n_z], without ell_factor for WL *)\n let kernel_bases = Array.make nt (Nx.zeros f64 [| n_z |]) in\n let kernel_has_ell_factor = Array.make nt false in\n Array.iteri\n (fun idx t ->\n match t.kind with\n | Weak_lensing { ia_bias; sigma_e = _; m_bias } ->\n let nz_tensor = nz_arrs.(idx) in\n (* A(z_j) = ∫_{z_j}^{zmax} n(z') dz' *)\n let a_vec = rev_cumtrapz nz_tensor n_z dz in\n (* B(z_j) = ∫_{z_j}^{zmax} n(z')/χ(z') dz' — tensor, through chi *)\n let nz_over_chi = Nx.div nz_tensor chi_safe in\n let b_vec = rev_cumtrapz nz_over_chi n_z dz in\n (* g = A - chi * B *)\n let g_vec = Nx.sub a_vec (Nx.mul chi_vec b_vec) in\n (* WL kernel base: (3 H0² Ωm / 2c) × (1+z) × χ × g *)\n let prefactor =\n Nx.mul_s omega_m_t (3.0 *. h0_ref *. h0_ref /. (2.0 *. c_km_s))\n in\n let one_plus_z = Nx.add_s z_vec 1.0 in\n let k_base =\n Nx.mul prefactor (Nx.mul one_plus_z (Nx.mul chi_vec g_vec))\n in\n (* Add NLA intrinsic alignment if present *)\n let k_base =\n match ia_bias with\n | None -> k_base\n | Some ia_b ->\n let ia_tensor =\n Nx.stack\n (List.init n_z (fun j -> ia_b p (Nx.scalar f64 z_arr.(j))))\n in\n (* K_IA = -(C₁ ρ_crit Ωm / D(z)) × n(z) × b_IA(z) × H(z) *)\n let ia_kernel =\n Nx.mul\n (Nx.mul_s omega_m_t (-.c1_rho_crit))\n (Nx.mul\n (Nx.div nz_tensor growth_vec)\n (Nx.mul ia_tensor h_vec))\n in\n Nx.add k_base ia_kernel\n in\n (* Shear multiplicative bias: W_obs = (1+m) W_true *)\n let k_base =\n if m_bias = 0.0 then k_base else Nx.mul_s k_base (1.0 +. m_bias)\n in\n kernel_bases.(idx) <- k_base;\n kernel_has_ell_factor.(idx) <- true\n | Number_counts { bias } ->\n let nz_tensor = nz_arrs.(idx) in\n let bias_tensor =\n Nx.stack (List.init n_z (fun j -> bias p (Nx.scalar f64 z_arr.(j))))\n in\n (* NC kernel: n(z) × b(z) × H(z) — no ell factor *)\n kernel_bases.(idx) <- Nx.mul nz_tensor (Nx.mul bias_tensor h_vec);\n kernel_has_ell_factor.(idx) <- false\n | Custom { kernel } ->\n (* Custom kernel: user provides the full W(z) *)\n kernel_bases.(idx) <-\n Nx.stack\n (List.init n_z (fun j ->\n let z_t = Nx.scalar f64 z_arr.(j) in\n let chi_t = Nx.get [ j ] chi_safe in\n kernel ~p ~z:z_t ~chi:chi_t));\n kernel_has_ell_factor.(idx) <- false)\n tracers_arr;\n\n (* Common integration weight: dchi/dz / chi² / c² × simpson *)\n let integ_weight =\n Nx.mul simpson_w\n (Nx.div_s\n (Nx.div dchi_dz_vec (Nx.mul chi_safe chi_safe))\n (c_km_s *. c_km_s))\n in\n\n (* Power spectrum grid [n_z, n_ell]: loop over z (scalar), vectorized over k.\n Both linear_power and nonlinear_power accept vector k but scalar z. *)\n let pk_grid =\n Nx.stack\n (List.init n_z (fun z_idx ->\n let z_t = Nx.scalar f64 z_arr.(z_idx) in\n let chi_z = Nx.get [ z_idx ] chi_safe in\n let k_vec = Nx.div (Nx.add_s ell 0.5) chi_z in\n power p k_vec z_t))\n in\n\n (* ell_factor vector [n_ell]: sqrt((ℓ-1)ℓ(ℓ+1)(ℓ+2)) / (ℓ+0.5)² *)\n let ell_factor_vec =\n let l = ell in\n let num =\n Nx.mul\n (Nx.mul (Nx.sub_s l 1.0) l)\n (Nx.mul (Nx.add_s l 1.0) (Nx.add_s l 2.0))\n in\n let den = Nx.mul (Nx.add_s l 0.5) (Nx.add_s l 0.5) in\n Nx.div (Nx.sqrt (Nx.abs num)) den\n in\n\n (* Limber integration: functional, no in-place mutation. integ_weight is\n [n_z], pk_grid is [n_z, n_ell]. For each pair (i,j): C_ℓ = Σ_z K_i(z)\n K_j(z) P(k,z) w(z) kernel_bases are [n_z], broadcast with pk_grid [n_z,\n n_ell]. *)\n let w_pk = Nx.mul (Nx.reshape [| n_z; 1 |] integ_weight) pk_grid in\n let spectra =\n Nx.stack\n (List.map\n (fun (i, j) ->\n let ki = Nx.reshape [| n_z; 1 |] kernel_bases.(i) in\n let kj = Nx.reshape [| n_z; 1 |] kernel_bases.(j) in\n let integrand = Nx.mul (Nx.mul ki kj) w_pk in\n let cl_row = Nx.sum ~axes:[ 0 ] integrand in\n let ell_power =\n (if kernel_has_ell_factor.(i) then 1 else 0)\n + if kernel_has_ell_factor.(j) then 1 else 0\n in\n if ell_power = 0 then cl_row\n else if ell_power = 1 then Nx.mul ell_factor_vec cl_row\n else Nx.mul (Nx.mul ell_factor_vec ell_factor_vec) cl_row)\n (Array.to_list pairs_arr))\n in\n { ell; tracers = tracers_arr; spectra }\n\n(* Cls submodule *)\n\nmodule Cls = struct\n let get cls ~i ~j =\n let n = Array.length cls.tracers in\n if i < 0 || i >= n || j < 0 || j >= n then\n invalid_arg \"Survey.Cls.get: index out of range\";\n Nx.slice [ I (pair_index n i j) ] cls.spectra\n\n let ell cls = cls.ell\n let n_tracers cls = Array.length cls.tracers\n let to_tensor cls = cls.spectra\n\n let noise cls =\n let n_ell = (Nx.shape cls.ell).(0) in\n let nt = Array.length cls.tracers in\n let pairs = cl_pairs nt in\n let n_cls = List.length pairs in\n let result = Nx.zeros f64 [| n_cls; n_ell |] in\n let pair_idx = ref 0 in\n List.iter\n (fun (i, j) ->\n if i = j then begin\n let noise_val = cls.tracers.(i).noise in\n for l = 0 to n_ell - 1 do\n Nx.set_item [ !pair_idx; l ] noise_val result\n done\n end;\n incr pair_idx)\n pairs;\n result\n\n let gaussian_covariance ?(f_sky = 0.25) cls =\n let ell = cls.ell in\n let n_ell = (Nx.shape ell).(0) in\n let nt = Array.length cls.tracers in\n let pairs = cl_pairs nt in\n let n_cls = List.length pairs in\n let n = n_cls * n_ell in\n let cov = Nx.zeros f64 [| n; n |] in\n let cl_noise = noise cls in\n let cl_obs = Nx.add cls.spectra cl_noise in\n let pairs_arr = Array.of_list pairs in\n let find_pair a b = pair_index nt a b in\n (* Δℓ via finite differences *)\n let dell =\n Array.init n_ell (fun l ->\n if l = 0 then Nx.item [ 1 ] ell -. Nx.item [ 0 ] ell\n else if l = n_ell - 1 then Nx.item [ l ] ell -. Nx.item [ l - 1 ] ell\n else 0.5 *. (Nx.item [ l + 1 ] ell -. Nx.item [ l - 1 ] ell))\n in\n for p1 = 0 to n_cls - 1 do\n let i, j = pairs_arr.(p1) in\n for p2 = p1 to n_cls - 1 do\n let m, nn = pairs_arr.(p2) in\n let im = find_pair i m and jn = find_pair j nn in\n let in_ = find_pair i nn and jm = find_pair j m in\n for l = 0 to n_ell - 1 do\n let ell_l = Nx.item [ l ] ell in\n let norm = ((2.0 *. ell_l) +. 1.0) *. dell.(l) *. f_sky in\n let c_im = Nx.get [ im; l ] cl_obs in\n let c_jn = Nx.get [ jn; l ] cl_obs in\n let c_in = Nx.get [ in_; l ] cl_obs in\n let c_jm = Nx.get [ jm; l ] cl_obs in\n let val_ =\n Nx.div_s (Nx.add (Nx.mul c_im c_jn) (Nx.mul c_in c_jm)) norm\n in\n let row = (p1 * n_ell) + l in\n let col = (p2 * n_ell) + l in\n Nx.set [ row; col ] cov val_;\n if p1 <> p2 then Nx.set [ col; row ] cov val_\n done\n done\n done;\n cov\nend\n"
},
{
"path": "dev/umbra/lib/survey.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Angular power spectra and survey science.\n\n The central type is {!tracer}: one tracer per tomographic bin. {!angular_cl}\n cross-correlates a list of tracers and returns a structured {!cls} value\n with typed accessors.\n\n {!angular_cl} and {!inverse_growth_bias} are differentiable through Rune.\n {!Cls.noise} and {!Cls.gaussian_covariance} are not (they use in-place\n mutation); compute them once at a fiducial cosmology.\n\n {[\n let nz1 = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.3 () in\n let nz2 = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.7 () in\n let wl1 = Survey.weak_lensing ~n_gal:26.0 nz1 in\n let wl2 = Survey.weak_lensing ~n_gal:26.0 nz2 in\n let ell = Nx.logspace Nx.float64 1.0 3.0 50 in\n let cls = Survey.angular_cl ~p:Cosmo.planck18 ~ell [ wl1; wl2 ] in\n let cl_auto = Survey.Cls.get cls ~i:0 ~j:0 in\n let cl_cross = Survey.Cls.get cls ~i:0 ~j:1\n ]} *)\n\n(** {1:nz Redshift distributions} *)\n\ntype nz\n(** A normalized redshift probability density n(z) with a maximum redshift. *)\n\nval smail : ?zmax:float -> a:float -> b:float -> z0:float -> unit -> nz\n(** [smail ~a ~b ~z0 ()] is n(z) {e ∝} z{^ a} exp(-(z/z0){^ b}). Auto-normalized\n via Simpson's rule. [zmax] defaults to [10.0]. *)\n\nval tabulated : z:Nx.float64_t -> pz:Nx.float64_t -> unit -> nz\n(** [tabulated ~z ~pz ()] is n(z) linearly interpolated from sampled points.\n Auto-normalized. [zmax] is inferred from the last element of [z]. *)\n\nval custom_nz : ?zmax:float -> (Nx.float64_t -> Nx.float64_t) -> nz\n(** [custom_nz f] is a redshift distribution with evaluation function [f]. [f z]\n maps a scalar tensor [z] to n(z). For differentiable survey optimization,\n [f] should use tensor operations so gradients flow through Rune. [zmax]\n defaults to [10.0]. *)\n\nval eval_nz : nz -> Nx.float64_t -> Nx.float64_t\n(** [eval_nz nz z] evaluates the normalized n(z) at [z]. *)\n\nval nz_zmax : nz -> float\n(** [nz_zmax nz] is the maximum redshift of the distribution. *)\n\n(** {1:bias Galaxy bias} *)\n\ntype bias = Cosmo.params -> Nx.float64_t -> Nx.float64_t\n(** A galaxy bias function. [bias p z] is b(z) under cosmology [p]. *)\n\nval constant_bias : float -> bias\n(** [constant_bias b] is a redshift-independent linear bias. Not differentiable\n (constant value). *)\n\nval inverse_growth_bias : float -> bias\n(** [inverse_growth_bias b0] is [b0 / D(z)], where D is the linear growth\n factor. Differentiable through Rune. *)\n\n(** {1:power Power spectrum backends} *)\n\ntype power = Cosmo.params -> Nx.float64_t -> Nx.float64_t -> Nx.float64_t\n(** [power p k z] is the matter power spectrum P(k, z) in (Mpc/h){^ 3}. [k] is a\n 1-D tensor of wavenumbers in h/Mpc, [z] is a scalar tensor. *)\n\nval linear : power\n(** [linear] is the linear matter power spectrum via Eisenstein & Hu (1998).\n Differentiable through Rune. *)\n\nval nonlinear : power\n(** [nonlinear] is the nonlinear power spectrum via Halofit (Takahashi et al.\n 2012). Differentiable through Rune (except the nonlinear scale k{_ nl} which\n is found by float-level root-finding). *)\n\nval baryonic_feedback :\n ?a_bary:float -> ?log10_k_star:float -> ?sigma:float -> power -> power\n(** [baryonic_feedback base_power] wraps [base_power] with a Gaussian\n suppression in log{_ 10}(k) that models baryonic feedback on the matter\n power spectrum:\n\n P{_ bary}(k, z) = P(k, z) {e ×} (1 - a{_ bary} {e ×} exp(-(log{_ 10}(k) -\n log{_ 10}(k{_ star})){^ 2} / {e σ}{^ 2})).\n\n [a_bary] is the suppression amplitude (default [0.0] = no effect).\n [log10_k_star] is the log{_ 10} of the peak suppression wavenumber in h/Mpc\n (default [1.0], i.e. k{_ star} = 10 h/Mpc). [sigma] is the Gaussian width in\n log{_ 10}(k) (default [0.55]).\n\n Differentiable through Rune. *)\n\n(** {1:tracers Tracers} *)\n\ntype tracer\n(** The type for a single tomographic tracer. One tracer = one redshift bin with\n its physics (lensing kernel, galaxy bias, etc.) and noise properties.\n {!angular_cl} cross-correlates a list of tracers. *)\n\nval weak_lensing :\n ?ia_bias:bias ->\n ?sigma_e:float ->\n ?m_bias:float ->\n ?n_gal:float ->\n nz ->\n tracer\n(** [weak_lensing nz] is a weak gravitational lensing tracer with redshift\n distribution [nz]. [sigma_e] is the intrinsic ellipticity dispersion\n (default [0.26]). [n_gal] is the galaxy number density in\n galaxies/arcmin{^ 2} (default [1.0]). [ia_bias], if provided, adds NLA\n intrinsic alignment.\n\n [m_bias] is the shear multiplicative bias (default [0.0]). The lensing\n kernel is scaled by [(1 + m_bias)], so auto-spectra scale as [(1 + m){^ 2}]\n and cross-spectra as [(1 + m{_ i})(1 + m{_ j})]. Differentiable through Rune\n when used with {!angular_cl}. *)\n\nval number_counts : bias:bias -> ?n_gal:float -> nz -> tracer\n(** [number_counts ~bias nz] is a galaxy number counts tracer with redshift\n distribution [nz] and galaxy bias model [bias]. [n_gal] is the galaxy number\n density in galaxies/arcmin{^ 2} (default [1.0]). *)\n\nval tracer :\n ?noise:float ->\n ?zmax:float ->\n (p:Cosmo.params -> z:Nx.float64_t -> chi:Nx.float64_t -> Nx.float64_t) ->\n tracer\n(** [tracer kernel] is a custom tracer with kernel function [kernel].\n\n [kernel ~p ~z ~chi] returns the full projection kernel W(z) at scalar\n redshift [z] and comoving distance [chi] (Mpc/h) under cosmology [p].\n\n [noise] is the constant noise power N{_ ℓ} for auto-correlations (default\n [0.0]). [zmax] defaults to [3.0]. *)\n\n(** {1:cls Angular power spectra} *)\n\ntype cls\n(** The type for a set of angular power spectra. Stores all auto- and\n cross-correlations for a list of tracers, along with the ell values and\n tracer metadata needed for noise and covariance computation. *)\n\nval angular_cl :\n ?p:Cosmo.params -> ?power:power -> ell:Nx.float64_t -> tracer list -> cls\n(** [angular_cl ~ell tracers] computes angular power spectra C{_ ℓ} for all\n auto- and cross-correlations via the Limber approximation. Differentiable\n through Rune.\n\n [power] defaults to {!nonlinear}. [p] defaults to {!Cosmo.planck18}.\n\n Raises [Invalid_argument] if [omega_b], [n_s], or [sigma8] are not set in\n [p]. *)\n\n(** {2:cls_access Structured access} *)\n\nmodule Cls : sig\n val get : cls -> i:int -> j:int -> Nx.float64_t\n (** [get cls ~i ~j] is the angular power spectrum C{_ ℓ}{^ ij} between tracers\n [i] and [j]. Returns a 1-D tensor of shape [[n_ell]]. [get cls ~i ~j] and\n [get cls ~j ~i] return the same spectrum.\n\n Raises [Invalid_argument] if [i] or [j] is out of range. *)\n\n val ell : cls -> Nx.float64_t\n (** [ell cls] is the multipole values, shape [[n_ell]]. *)\n\n val n_tracers : cls -> int\n (** [n_tracers cls] is the number of tracers. *)\n\n val to_tensor : cls -> Nx.float64_t\n (** [to_tensor cls] is all spectra packed as a tensor of shape\n [[n_cls; n_ell]] where [n_cls = n * (n + 1) / 2], ordered as (0,0), (0,1),\n ..., (1,1), .... *)\n\n val noise : cls -> Nx.float64_t\n (** [noise cls] is the shot noise power spectra. Weak lensing:\n {e σ}{_ e}{^ 2}/n{_ gal}. Number counts: 1/n{_ gal}. Custom: the [noise]\n value. Cross-spectra are zero. Shape [[n_cls; n_ell]].\n\n Not differentiable. *)\n\n val gaussian_covariance : ?f_sky:float -> cls -> Nx.float64_t\n (** [gaussian_covariance cls] is the Gaussian covariance matrix. [f_sky]\n defaults to [0.25]. Returns dense matrix of shape [[n; n]] where\n [n = n_cls * n_ell].\n\n Not differentiable. *)\nend\n"
},
{
"path": "dev/umbra/lib/time.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Astronomical time with phantom-typed time scales.\n\n Internal representation: Julian Date (float) in the tagged scale. MJD = JD -\n 2400000.5 Unix epoch (1970-01-01T00:00:00 UTC) = JD 2440587.5 *)\n\ntype 'a t = float\ntype utc\ntype tai\ntype tt\ntype tdb\n\n(* Constructors *)\n\nlet unsafe_of_jd jd = jd\nlet unsafe_of_mjd mjd = mjd +. 2_400_000.5\nlet of_unix u = (u /. 86_400.0) +. 2_440_587.5\nlet now () = of_unix (Unix.gettimeofday ())\n\n(* Comparison *)\n\nlet compare (a : float) (b : float) = Float.compare a b\nlet equal (a : float) (b : float) = Float.equal a b\n\n(* Eliminators *)\n\nlet to_jd t = t\nlet to_mjd t = t -. 2_400_000.5\nlet to_unix t = (t -. 2_440_587.5) *. 86_400.0\n\n(* Duration *)\n\nlet diff a b = Unit.Time.day (a -. b)\nlet add t dt = t +. Nx.item [] (Unit.Time.in_day dt)\n\n(* Leap second table: (JD of midnight UTC when leap second is introduced,\n cumulative TAI-UTC). Source: IERS Bulletin C. *)\nlet leap_seconds =\n [|\n (2441317.5, 10.0);\n (* 1972-01-01 *)\n (2441499.5, 11.0);\n (* 1972-07-01 *)\n (2441683.5, 12.0);\n (* 1973-01-01 *)\n (2442048.5, 13.0);\n (* 1974-01-01 *)\n (2442413.5, 14.0);\n (* 1975-01-01 *)\n (2442778.5, 15.0);\n (* 1976-01-01 *)\n (2443144.5, 16.0);\n (* 1977-01-01 *)\n (2443509.5, 17.0);\n (* 1978-01-01 *)\n (2443874.5, 18.0);\n (* 1979-01-01 *)\n (2444239.5, 19.0);\n (* 1980-01-01 *)\n (2444786.5, 20.0);\n (* 1981-07-01 *)\n (2445151.5, 21.0);\n (* 1982-07-01 *)\n (2445516.5, 22.0);\n (* 1983-07-01 *)\n (2446247.5, 23.0);\n (* 1985-07-01 *)\n (2447161.5, 24.0);\n (* 1988-01-01 *)\n (2447892.5, 25.0);\n (* 1990-01-01 *)\n (2448257.5, 26.0);\n (* 1991-01-01 *)\n (2448804.5, 27.0);\n (* 1992-07-01 *)\n (2449169.5, 28.0);\n (* 1993-07-01 *)\n (2449534.5, 29.0);\n (* 1994-07-01 *)\n (2450083.5, 30.0);\n (* 1996-01-01 *)\n (2450630.5, 31.0);\n (* 1997-07-01 *)\n (2451179.5, 32.0);\n (* 1999-01-01 *)\n (2453736.5, 33.0);\n (* 2006-01-01 *)\n (2454832.5, 34.0);\n (* 2009-01-01 *)\n (2456109.5, 35.0);\n (* 2012-07-01 *)\n (2457204.5, 36.0);\n (* 2015-07-01 *)\n (2457754.5, 37.0);\n (* 2017-01-01 *)\n |]\n\nlet tai_minus_utc jd_utc =\n let n = Array.length leap_seconds in\n let rec search i =\n if i < 0 then 10.0\n else\n let jd, dt = leap_seconds.(i) in\n if jd_utc >= jd then dt else search (i - 1)\n in\n search (n - 1)\n\n(* UTC <-> TAI *)\n\nlet utc_to_tai utc_jd =\n let dt = tai_minus_utc utc_jd in\n utc_jd +. (dt /. 86_400.0)\n\nlet tai_to_utc tai_jd =\n (* Approximate: convert TAI to approximate UTC, look up, refine *)\n let approx_utc = tai_jd -. (37.0 /. 86_400.0) in\n let dt = tai_minus_utc approx_utc in\n tai_jd -. (dt /. 86_400.0)\n\n(* TAI <-> TT: TT = TAI + 32.184s (exact by definition) *)\n\nlet tt_offset = 32.184 /. 86_400.0\nlet tai_to_tt tai_jd = tai_jd +. tt_offset\nlet tt_to_tai tt_jd = tt_jd -. tt_offset\n\n(* TT <-> TDB: Fairhead & Bretagnon 1990 series (first 10 terms). Accuracy ~1μs\n for dates within a few centuries of J2000.0.\n\n T = (JD_TT - 2451545.0) / 36525.0 (Julian centuries from J2000.0 TT) TDB - TT\n ≈ Σ Aᵢ sin(ωᵢ T + φᵢ) in seconds *)\n\nlet fb_terms =\n [|\n (* amplitude (s), frequency (rad/century), phase (rad) *)\n (1.656_674_564e-3, 6_283.075_849_991, 6.240_054_195);\n (2.227_2e-5, 5_753.384_884_897, 4.296_977_442);\n (1.3886e-5, 12_566.151_699_983, 6.196_904_410);\n (3.150e-6, 529.690_965_095, 0.444_401_603);\n (1.575e-6, 6_069.776_754_553, 4.021_195_093);\n (1.020_5e-5, 213.299_095_438, 5.543_113_262);\n (3.978e-6, 77_713.771_467_920, 5.198_467_090);\n (4.354e-6, 7_860.419_392_439, 5.988_822_341);\n (1.456e-6, 11_506.769_769_794, 2.457_236_222);\n (1.126e-6, 3_930.209_696_220, 5.316_024_159);\n |]\n\nlet tt_to_tdb tt_jd =\n let t = (tt_jd -. 2_451_545.0) /. 36_525.0 in\n let sum = ref 0.0 in\n for i = 0 to Array.length fb_terms - 1 do\n let amp, freq, phase = fb_terms.(i) in\n sum := !sum +. (amp *. Float.sin ((freq *. t) +. phase))\n done;\n tt_jd +. (!sum /. 86_400.0)\n\nlet tdb_to_tt tdb_jd =\n (* Single Newton iteration: TT ≈ TDB, compute correction *)\n let tt_approx = tdb_jd in\n let tdb_from_approx = tt_to_tdb tt_approx in\n let correction = tdb_jd -. tdb_from_approx in\n tt_approx +. correction\n\n(* ISO 8601 parsing and formatting for UTC *)\n\n(* Calendar date to JD (valid for dates after 1582-10-15, Gregorian calendar) *)\nlet cal_to_jd y m d =\n let y, m = if m <= 2 then (y - 1, m + 12) else (y, m) in\n let a = y / 100 in\n let b = 2 - a + (a / 4) in\n Float.floor (365.25 *. Float.of_int (y + 4716))\n +. Float.floor (30.6001 *. Float.of_int (m + 1))\n +. d +. Float.of_int b -. 1524.5\n\n(* JD to calendar date *)\nlet jd_to_cal jd =\n let jd = jd +. 0.5 in\n let z = Float.to_int (Float.floor jd) in\n let f = jd -. Float.of_int z in\n let a =\n if z < 2299161 then z\n else\n let alpha =\n Float.to_int (Float.floor ((Float.of_int z -. 1867216.25) /. 36524.25))\n in\n z + 1 + alpha - (alpha / 4)\n in\n let b = a + 1524 in\n let c = Float.to_int (Float.floor ((Float.of_int b -. 122.1) /. 365.25)) in\n let d = Float.to_int (Float.floor (365.25 *. Float.of_int c)) in\n let e = Float.to_int (Float.floor (Float.of_int (b - d) /. 30.6001)) in\n let day_frac =\n Float.of_int (b - d) -. Float.floor (30.6001 *. Float.of_int e) +. f\n in\n let month = if e < 14 then e - 1 else e - 13 in\n let year = if month > 2 then c - 4716 else c - 4715 in\n (year, month, day_frac)\n\nlet of_iso s =\n let s =\n let len = String.length s in\n if len > 0 && s.[len - 1] = 'Z' then String.sub s 0 (len - 1) else s\n in\n match\n Scanf.sscanf s \"%d-%d-%dT%d:%d:%f\" (fun y mo d h mi s ->\n (y, mo, d, h, mi, s))\n with\n | y, mo, d, h, mi, sec ->\n let day =\n Float.of_int d\n +. (Float.of_int h /. 24.0)\n +. (Float.of_int mi /. 1440.0)\n +. (sec /. 86_400.0)\n in\n cal_to_jd y mo day\n | exception _ -> (\n match Scanf.sscanf s \"%d-%d-%d\" (fun y mo d -> (y, mo, d)) with\n | y, mo, d -> cal_to_jd y mo (Float.of_int d)\n | exception _ -> invalid_arg (\"Time.of_iso: cannot parse \" ^ s))\n\nlet to_iso t =\n let y, m, day_frac = jd_to_cal t in\n let d = Float.to_int (Float.floor day_frac) in\n let frac = day_frac -. Float.of_int d in\n let total_sec = frac *. 86_400.0 in\n let h = Float.to_int (Float.floor (total_sec /. 3600.0)) in\n let rem = total_sec -. (Float.of_int h *. 3600.0) in\n let mi = Float.to_int (Float.floor (rem /. 60.0)) in\n let sec = rem -. (Float.of_int mi *. 60.0) in\n if Float.abs sec < 0.0005 then\n Printf.sprintf \"%04d-%02d-%02dT%02d:%02d:%02dZ\" y m d h mi 0\n else Printf.sprintf \"%04d-%02d-%02dT%02d:%02d:%06.3fZ\" y m d h mi sec\n"
},
{
"path": "dev/umbra/lib/time.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Astronomical time with phantom-typed time scales.\n\n Times are stored internally as Julian Dates (float). Scale conversions are\n type-safe: {!utc_to_tai} accepts a [utc t] and returns a [tai t].\n\n {[\n let t = Time.of_iso \"2024-01-01T00:00:00\" in\n let tai = Time.utc_to_tai t in\n let tt = Time.tai_to_tt tai in\n let jd = Time.to_jd tt\n ]} *)\n\n(** {1:types Types} *)\n\ntype 'scale t\n(** The type for a Julian Date tagged with time scale ['scale]. *)\n\ntype utc\n(** Coordinated Universal Time. *)\n\ntype tai\n(** International Atomic Time. *)\n\ntype tt\n(** Terrestrial Time. *)\n\ntype tdb\n(** Barycentric Dynamical Time. *)\n\n(** {1:constructors Constructors} *)\n\nval unsafe_of_jd : float -> 'a t\n(** [unsafe_of_jd jd] is a time from the Julian Date [jd]. The caller must\n ensure [jd] is in the intended time scale. *)\n\nval unsafe_of_mjd : float -> 'a t\n(** [unsafe_of_mjd mjd] is a time from the Modified Julian Date [mjd] (MJD = JD\n \\- 2400000.5). The caller must ensure [mjd] is in the intended time scale.\n*)\n\nval of_iso : string -> utc t\n(** [of_iso s] parses an ISO 8601 date-time string as UTC. Accepted formats:\n [\"YYYY-MM-DD\"], [\"YYYY-MM-DDThh:mm:ss\"], and [\"YYYY-MM-DDThh:mm:ssZ\"].\n\n {b Warning.} Uses the Gregorian calendar; dates before 1582-10-15 produce\n incorrect Julian Dates. Leap seconds (e.g. [23:59:60]) cannot be represented\n and are parsed as the following second.\n\n Raises [Invalid_argument] if [s] cannot be parsed. *)\n\nval of_unix : float -> utc t\n(** [of_unix u] is the UTC time corresponding to the Unix timestamp [u] (seconds\n since 1970-01-01T00:00:00 UTC). *)\n\nval now : unit -> utc t\n(** [now ()] is the current UTC time from the system clock. *)\n\n(** {1:comparison Comparison} *)\n\nval compare : 'a t -> 'a t -> int\n(** [compare a b] orders times by their Julian Date values. *)\n\nval equal : 'a t -> 'a t -> bool\n(** [equal a b] is [true] iff [a] and [b] have the same Julian Date value. *)\n\n(** {1:eliminators Eliminators} *)\n\nval to_jd : 'a t -> float\n(** [to_jd t] is the Julian Date of [t]. *)\n\nval to_mjd : 'a t -> float\n(** [to_mjd t] is the Modified Julian Date of [t] (MJD = JD - 2400000.5). *)\n\nval to_iso : utc t -> string\n(** [to_iso t] formats [t] as an ISO 8601 string with trailing [Z]. Output is\n [\"YYYY-MM-DDThh:mm:ssZ\"] when the fractional seconds are below 0.5 ms, or\n [\"YYYY-MM-DDThh:mm:ss.sssZ\"] otherwise.\n\n {b Warning.} Leap-second labels like [23:59:60] cannot be produced; times\n within a leap second round to [00:00:00] of the following day. *)\n\nval to_unix : utc t -> float\n(** [to_unix t] is the Unix timestamp of [t] (seconds since 1970-01-01T00:00:00\n UTC). *)\n\n(** {1:scales Scale conversions}\n\n UTC/TAI conversions use the IERS leap-second table (Bulletin C), currently\n covering 1972-01-01 through 2017-01-01 (TAI-UTC = 37 s). Dates before\n 1972-01-01 use TAI-UTC = 10 s.\n\n TT = TAI + 32.184 s (exact by definition).\n\n TDB-TT uses the first 10 terms of the Fairhead & Bretagnon (1990) series,\n accurate to ~1 us within a few centuries of J2000.0. *)\n\nval utc_to_tai : utc t -> tai t\n(** [utc_to_tai t] converts [t] from UTC to TAI. *)\n\nval tai_to_utc : tai t -> utc t\n(** [tai_to_utc t] converts [t] from TAI to UTC. *)\n\nval tai_to_tt : tai t -> tt t\n(** [tai_to_tt t] converts [t] from TAI to TT. *)\n\nval tt_to_tai : tt t -> tai t\n(** [tt_to_tai t] converts [t] from TT to TAI. *)\n\nval tt_to_tdb : tt t -> tdb t\n(** [tt_to_tdb t] converts [t] from TT to TDB. *)\n\nval tdb_to_tt : tdb t -> tt t\n(** [tdb_to_tt t] converts [t] from TDB to TT. Uses a single Newton iteration;\n accurate to ~1 us. *)\n\n(** {1:duration Duration} *)\n\nval diff : 'a t -> 'a t -> Unit.time Unit.t\n(** [diff a b] is the duration [a - b] as a {!Unit.time} quantity. *)\n\nval add : 'a t -> Unit.time Unit.t -> 'a t\n(** [add t dt] is [t] offset by the duration [dt]. *)\n"
},
{
"path": "dev/umbra/lib/umbra.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Unit = Unit\nmodule Const = Const\nmodule Time = Time\nmodule Coord = Coord\nmodule Altaz = Altaz\nmodule Galactocentric = Galactocentric\nmodule Cosmo = Cosmo\nmodule Spectrum = Spectrum\nmodule Extinction = Extinction\nmodule Photometry = Photometry\nmodule Filters = Filters\nmodule Survey = Survey\n"
},
{
"path": "dev/umbra/lib/umbra.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Computational astronomy for OCaml.\n\n Umbra provides dimensionally-typed physical quantities, astronomical\n constants, cosmological distance calculations, spectral energy\n distributions, dust extinction, synthetic photometry, coordinate transforms,\n time scales, and catalog cross-matching.\n\n All computations operate on {!Nx} tensors and are differentiable through\n {!Rune} by default.\n\n {[\n open Umbra\n\n let z = Nx.scalar Nx.float64 0.5 in\n let dl = Cosmo.luminosity_distance z in\n let dl_mpc = Unit.Length.in_mpc dl in\n\n let rv = Nx.scalar Nx.float64 3.1 in\n let av = Nx.scalar Nx.float64 0.5 in\n let wave = Unit.Length.of_m (Nx.linspace Nx.float64 3e-7 1e-6 1000) in\n let bp = Photometry.tophat\n ~lo:(Unit.Length.nm 400.0) ~hi:(Unit.Length.nm 700.0) ~n:1000 in\n let sed =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar Nx.float64 5800.0))\n ~wavelength:wave\n |> Extinction.apply (Extinction.ccm89 ~rv) ~av\n |> Spectrum.as_flux_density\n in\n let mag = Photometry.ab_mag bp sed\n ]} *)\n\n(** {1:units Units and constants} *)\n\nmodule Unit = Unit\n(** Physical quantities with compile-time dimensional safety. *)\n\nmodule Const = Const\n(** Physical and astronomical constants (CODATA 2022, IAU 2015). *)\n\n(** {1:astro Astronomy} *)\n\nmodule Time = Time\n(** Astronomical time with phantom-typed time scales. *)\n\nmodule Coord = Coord\n(** Celestial coordinates with frame transforms and catalog cross-matching. *)\n\nmodule Altaz = Altaz\n(** Altitude-azimuth (horizontal) coordinates. *)\n\nmodule Galactocentric = Galactocentric\n(** Galactocentric Cartesian coordinates. *)\n\nmodule Cosmo = Cosmo\n(** Cosmological distances for {e Λ}CDM, wCDM, and w0waCDM universes. *)\n\nmodule Spectrum = Spectrum\n(** Sampled spectral values on a wavelength grid. *)\n\nmodule Extinction = Extinction\n(** Dust extinction laws. *)\n\nmodule Photometry = Photometry\n(** Synthetic photometry over filter bandpasses. *)\n\nmodule Filters = Filters\n(** Standard astronomical filter bandpasses (SDSS, Johnson-Cousins, 2MASS, Gaia\n DR3). *)\n\n(** {1:survey Survey science} *)\n\nmodule Survey = Survey\n(** Angular power spectra, probes, and survey likelihood. *)\n"
},
{
"path": "dev/umbra/lib/unit.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet f64 = Nx.float64\n\ntype 'a t = Nx.float64_t\ntype length\ntype mass\ntype time\ntype angle\ntype velocity\ntype power\ntype temperature\ntype energy\ntype frequency\ntype dimensionless\n\n(* Arithmetic — all Nx ops, fully traced by rune *)\n\nlet ( + ) a b = Nx.add a b\nlet ( - ) a b = Nx.sub a b\nlet neg x = Nx.neg x\nlet abs x = Nx.abs x\nlet scale s x = Nx.mul_s x s\nlet scale_t s x = Nx.mul s x\nlet ratio a b = Nx.div a b\nlet zero = Nx.scalar f64 0.0\nlet compare a b = Float.compare (Nx.item [] a) (Nx.item [] b)\nlet equal a b = Float.equal (Nx.item [] a) (Nx.item [] b)\nlet pp fmt x = Format.fprintf fmt \"%g\" (Nx.item [] x)\nlet to_float x = Nx.item [] x\n\n(* Physical constants used in cross-dimension combinators *)\n\nlet c_m_s = Nx.scalar f64 299_792_458.0\nlet h_si = Nx.scalar f64 6.626_070_15e-34\nlet hc_si = Nx.scalar f64 (6.626_070_15e-34 *. 299_792_458.0)\nlet one = Nx.scalar f64 1.0\nlet au_m_t = Nx.scalar f64 1.495_978_707e11\n\n(* Cross-dimension combinators — all Nx ops *)\n\nlet length_per_time d t = Nx.div d t\nlet velocity_times_time v t = Nx.mul v t\nlet length_per_velocity d v = Nx.div d v\nlet wavelength_to_frequency lam = Nx.div c_m_s lam\nlet frequency_to_wavelength nu = Nx.div c_m_s nu\nlet frequency_to_energy nu = Nx.mul h_si nu\nlet energy_to_frequency e = Nx.div e h_si\nlet energy_to_wavelength e = Nx.div hc_si e\n\n(* Parallax: 1 arcsec ↔ 1 parsec. parallax(rad) = 1 AU / distance(m), so\n distance(m) = 1 AU / parallax(rad). Uses the scale factors defined below. *)\n\nlet parallax_to_distance p = Nx.div au_m_t p\nlet distance_to_parallax d = Nx.div au_m_t d\n\n(* Spectral density: f_ν = f_λ · λ²/c, where f_λ is per-metre and f_ν is\n per-hertz. *)\nlet flam_to_fnu ~wavelength flam =\n Nx.div (Nx.mul flam (Nx.square wavelength)) c_m_s\n\nlet fnu_to_flam ~wavelength fnu =\n Nx.div (Nx.mul fnu c_m_s) (Nx.square wavelength)\n\n(* Doppler conventions: velocity ↔ observed wavelength given rest wavelength.\n All three conventions agree at v << c. *)\n\n(* Optical: z = v/c, λ_obs = λ_rest * (1 + v/c) *)\nlet doppler_optical ~rest v = Nx.mul rest (Nx.add_s (Nx.div v c_m_s) 1.0)\n\nlet doppler_optical_inv ~rest obs =\n Nx.mul c_m_s (Nx.sub_s (Nx.div obs rest) 1.0)\n\n(* Radio: v = c*(1 - λ_rest/λ_obs), λ_obs = λ_rest / (1 - v/c) *)\nlet doppler_radio ~rest v = Nx.div rest (Nx.sub one (Nx.div v c_m_s))\nlet doppler_radio_inv ~rest obs = Nx.mul c_m_s (Nx.sub one (Nx.div rest obs))\n\n(* Relativistic: λ_obs = λ_rest * sqrt((1+β)/(1-β)), β = v/c *)\nlet doppler_relativistic ~rest v =\n let beta = Nx.div v c_m_s in\n Nx.mul rest (Nx.sqrt (Nx.div (Nx.add_s beta 1.0) (Nx.sub one beta)))\n\nlet doppler_relativistic_inv ~rest obs =\n let r2 = Nx.square (Nx.div obs rest) in\n Nx.mul c_m_s (Nx.div (Nx.sub_s r2 1.0) (Nx.add_s r2 1.0))\n\n(* Scale factors to SI base unit *)\n\nlet pc_m = 3.085_677_581_491_367_3e16\nlet au_m = 1.495_978_707e11\nlet ly_m = 9.460_730_472_580_8e15\nlet solar_radius_m = 6.957e8\nlet earth_radius_m = 6.371e6\nlet jupiter_radius_m = 7.1492e7\nlet solar_mass_kg = 1.988_4e30\nlet earth_mass_kg = 5.972_2e24\nlet jupiter_mass_kg = 1.898_2e27\nlet solar_luminosity_w = 3.828e26\nlet julian_year_s = 365.25 *. 86_400.0\nlet ev_j = 1.602_176_634e-19\n\nmodule Length = struct\n let of_tensor x = x\n let to_tensor x = x\n let m x = Nx.scalar f64 x\n let km x = Nx.scalar f64 (x *. 1e3)\n let cm x = Nx.scalar f64 (x *. 1e-2)\n let mm x = Nx.scalar f64 (x *. 1e-3)\n let um x = Nx.scalar f64 (x *. 1e-6)\n let nm x = Nx.scalar f64 (x *. 1e-9)\n let angstrom x = Nx.scalar f64 (x *. 1e-10)\n let au x = Nx.scalar f64 (x *. au_m)\n let pc x = Nx.scalar f64 (x *. pc_m)\n let kpc x = Nx.scalar f64 (x *. pc_m *. 1e3)\n let mpc x = Nx.scalar f64 (x *. pc_m *. 1e6)\n let gpc x = Nx.scalar f64 (x *. pc_m *. 1e9)\n let ly x = Nx.scalar f64 (x *. ly_m)\n let solar_radius x = Nx.scalar f64 (x *. solar_radius_m)\n let earth_radius x = Nx.scalar f64 (x *. earth_radius_m)\n let jupiter_radius x = Nx.scalar f64 (x *. jupiter_radius_m)\n let of_m x = x\n let of_km x = Nx.mul_s x 1e3\n let of_cm x = Nx.mul_s x 1e-2\n let of_mm x = Nx.mul_s x 1e-3\n let of_um x = Nx.mul_s x 1e-6\n let of_nm x = Nx.mul_s x 1e-9\n let of_angstrom x = Nx.mul_s x 1e-10\n let of_au x = Nx.mul_s x au_m\n let of_pc x = Nx.mul_s x pc_m\n let of_kpc x = Nx.mul_s x (pc_m *. 1e3)\n let of_mpc x = Nx.mul_s x (pc_m *. 1e6)\n let of_gpc x = Nx.mul_s x (pc_m *. 1e9)\n let of_ly x = Nx.mul_s x ly_m\n let of_solar_radius x = Nx.mul_s x solar_radius_m\n let of_earth_radius x = Nx.mul_s x earth_radius_m\n let of_jupiter_radius x = Nx.mul_s x jupiter_radius_m\n let in_m x = x\n let in_km x = Nx.div_s x 1e3\n let in_cm x = Nx.mul_s x (1.0 /. 1e-2)\n let in_mm x = Nx.mul_s x (1.0 /. 1e-3)\n let in_um x = Nx.mul_s x (1.0 /. 1e-6)\n let in_nm x = Nx.mul_s x (1.0 /. 1e-9)\n let in_angstrom x = Nx.mul_s x (1.0 /. 1e-10)\n let in_au x = Nx.div_s x au_m\n let in_pc x = Nx.div_s x pc_m\n let in_kpc x = Nx.div_s x (pc_m *. 1e3)\n let in_mpc x = Nx.div_s x (pc_m *. 1e6)\n let in_gpc x = Nx.div_s x (pc_m *. 1e9)\n let in_ly x = Nx.div_s x ly_m\n let in_solar_radius x = Nx.div_s x solar_radius_m\n let in_earth_radius x = Nx.div_s x earth_radius_m\n let in_jupiter_radius x = Nx.div_s x jupiter_radius_m\nend\n\nmodule Mass = struct\n let of_tensor x = x\n let to_tensor x = x\n let kg x = Nx.scalar f64 x\n let g x = Nx.scalar f64 (x *. 1e-3)\n let mg x = Nx.scalar f64 (x *. 1e-6)\n let solar_mass x = Nx.scalar f64 (x *. solar_mass_kg)\n let earth_mass x = Nx.scalar f64 (x *. earth_mass_kg)\n let jupiter_mass x = Nx.scalar f64 (x *. jupiter_mass_kg)\n let of_kg x = x\n let of_g x = Nx.mul_s x 1e-3\n let of_mg x = Nx.mul_s x 1e-6\n let of_solar_mass x = Nx.mul_s x solar_mass_kg\n let of_earth_mass x = Nx.mul_s x earth_mass_kg\n let of_jupiter_mass x = Nx.mul_s x jupiter_mass_kg\n let in_kg x = x\n let in_g x = Nx.mul_s x (1.0 /. 1e-3)\n let in_mg x = Nx.mul_s x (1.0 /. 1e-6)\n let in_solar_mass x = Nx.div_s x solar_mass_kg\n let in_earth_mass x = Nx.div_s x earth_mass_kg\n let in_jupiter_mass x = Nx.div_s x jupiter_mass_kg\nend\n\nmodule Time = struct\n let of_tensor x = x\n let to_tensor x = x\n let s x = Nx.scalar f64 x\n let ms x = Nx.scalar f64 (x *. 1e-3)\n let us x = Nx.scalar f64 (x *. 1e-6)\n let min x = Nx.scalar f64 (x *. 60.0)\n let hr x = Nx.scalar f64 (x *. 3600.0)\n let day x = Nx.scalar f64 (x *. 86_400.0)\n let yr x = Nx.scalar f64 (x *. julian_year_s)\n let myr x = Nx.scalar f64 (x *. julian_year_s *. 1e6)\n let gyr x = Nx.scalar f64 (x *. julian_year_s *. 1e9)\n let of_s x = x\n let of_ms x = Nx.mul_s x 1e-3\n let of_us x = Nx.mul_s x 1e-6\n let of_min x = Nx.mul_s x 60.0\n let of_hr x = Nx.mul_s x 3600.0\n let of_day x = Nx.mul_s x 86_400.0\n let of_yr x = Nx.mul_s x julian_year_s\n let of_myr x = Nx.mul_s x (julian_year_s *. 1e6)\n let of_gyr x = Nx.mul_s x (julian_year_s *. 1e9)\n let in_s x = x\n let in_ms x = Nx.mul_s x (1.0 /. 1e-3)\n let in_us x = Nx.mul_s x (1.0 /. 1e-6)\n let in_min x = Nx.div_s x 60.0\n let in_hr x = Nx.div_s x 3600.0\n let in_day x = Nx.div_s x 86_400.0\n let in_yr x = Nx.div_s x julian_year_s\n let in_myr x = Nx.div_s x (julian_year_s *. 1e6)\n let in_gyr x = Nx.div_s x (julian_year_s *. 1e9)\nend\n\nmodule Angle = struct\n let deg_rad = Float.pi /. 180.0\n let of_tensor x = x\n let to_tensor x = x\n let rad x = Nx.scalar f64 x\n let deg x = Nx.scalar f64 (x *. deg_rad)\n let arcmin x = Nx.scalar f64 (x *. deg_rad /. 60.0)\n let arcsec x = Nx.scalar f64 (x *. deg_rad /. 3600.0)\n let mas x = Nx.scalar f64 (x *. deg_rad /. 3_600_000.0)\n let hour_angle x = Nx.scalar f64 (x *. Float.pi /. 12.0)\n let of_rad x = x\n let of_deg x = Nx.mul_s x deg_rad\n let of_arcmin x = Nx.mul_s x (deg_rad /. 60.0)\n let of_arcsec x = Nx.mul_s x (deg_rad /. 3600.0)\n let of_mas x = Nx.mul_s x (deg_rad /. 3_600_000.0)\n let of_hour_angle x = Nx.mul_s x (Float.pi /. 12.0)\n let in_rad x = x\n let in_deg x = Nx.div_s x deg_rad\n let in_arcmin x = Nx.mul_s (Nx.div_s x deg_rad) 60.0\n let in_arcsec x = Nx.mul_s (Nx.div_s x deg_rad) 3600.0\n let in_mas x = Nx.mul_s (Nx.div_s x deg_rad) 3_600_000.0\n let in_hour_angle x = Nx.mul_s x (12.0 /. Float.pi)\n let sin x = Nx.sin x\n let cos x = Nx.cos x\n let tan x = Nx.tan x\n let asin x = Nx.asin x\n let acos x = Nx.acos x\n let atan2 ~y ~x = Nx.atan2 y x\n\n let wrap_360 x =\n let d = in_deg x in\n let d = Nx.sub d (Nx.mul_s (Nx.floor (Nx.div_s d 360.0)) 360.0) in\n of_deg d\n\n let wrap_180 x =\n let d = Nx.add_s (in_deg x) 180.0 in\n let d = Nx.sub d (Nx.mul_s (Nx.floor (Nx.div_s d 360.0)) 360.0) in\n of_deg (Nx.sub_s d 180.0)\nend\n\nmodule Velocity = struct\n let of_tensor x = x\n let to_tensor x = x\n let m_s x = Nx.scalar f64 x\n let km_s x = Nx.scalar f64 (x *. 1e3)\n let km_hr x = Nx.scalar f64 (x *. (1e3 /. 3600.0))\n let of_m_s x = x\n let of_km_s x = Nx.mul_s x 1e3\n let of_km_hr x = Nx.mul_s x (1e3 /. 3600.0)\n let in_m_s x = x\n let in_km_s x = Nx.div_s x 1e3\n let in_km_hr x = Nx.div_s x (1e3 /. 3600.0)\nend\n\nmodule Power = struct\n let of_tensor x = x\n let to_tensor x = x\n let w x = Nx.scalar f64 x\n let kw x = Nx.scalar f64 (x *. 1e3)\n let solar_luminosity x = Nx.scalar f64 (x *. solar_luminosity_w)\n let erg_s x = Nx.scalar f64 (x *. 1e-7)\n let of_w x = x\n let of_kw x = Nx.mul_s x 1e3\n let of_solar_luminosity x = Nx.mul_s x solar_luminosity_w\n let of_erg_s x = Nx.mul_s x 1e-7\n let in_w x = x\n let in_kw x = Nx.div_s x 1e3\n let in_solar_luminosity x = Nx.div_s x solar_luminosity_w\n let in_erg_s x = Nx.mul_s x (1.0 /. 1e-7)\nend\n\nmodule Temperature = struct\n let of_tensor x = x\n let to_tensor x = x\n let kelvin x = Nx.scalar f64 x\n let of_kelvin x = x\n let in_kelvin x = x\nend\n\nmodule Energy = struct\n let of_tensor x = x\n let to_tensor x = x\n let j x = Nx.scalar f64 x\n let erg x = Nx.scalar f64 (x *. 1e-7)\n let ev x = Nx.scalar f64 (x *. ev_j)\n let kev x = Nx.scalar f64 (x *. ev_j *. 1e3)\n let mev x = Nx.scalar f64 (x *. ev_j *. 1e6)\n let of_j x = x\n let of_erg x = Nx.mul_s x 1e-7\n let of_ev x = Nx.mul_s x ev_j\n let of_kev x = Nx.mul_s x (ev_j *. 1e3)\n let of_mev x = Nx.mul_s x (ev_j *. 1e6)\n let in_j x = x\n let in_erg x = Nx.mul_s x (1.0 /. 1e-7)\n let in_ev x = Nx.div_s x ev_j\n let in_kev x = Nx.div_s x (ev_j *. 1e3)\n let in_mev x = Nx.div_s x (ev_j *. 1e6)\nend\n\nmodule Frequency = struct\n let of_tensor x = x\n let to_tensor x = x\n let hz x = Nx.scalar f64 x\n let khz x = Nx.scalar f64 (x *. 1e3)\n let mhz x = Nx.scalar f64 (x *. 1e6)\n let ghz x = Nx.scalar f64 (x *. 1e9)\n let of_hz x = x\n let of_khz x = Nx.mul_s x 1e3\n let of_mhz x = Nx.mul_s x 1e6\n let of_ghz x = Nx.mul_s x 1e9\n let in_hz x = x\n let in_khz x = Nx.div_s x 1e3\n let in_mhz x = Nx.div_s x 1e6\n let in_ghz x = Nx.div_s x 1e9\nend\n\nmodule Dimensionless = struct\n let of_tensor x = x\n let to_tensor x = x\n let v x = Nx.scalar f64 x\n let to_float x = Nx.item [] x\nend\n"
},
{
"path": "dev/umbra/lib/unit.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Physical quantities with compile-time dimensional safety.\n\n A {e quantity} is an {!Nx.float64_t} tensor of arbitrary shape tagged with a\n phantom dimension type. Arithmetic requires matching dimensions:\n [length t + length t] typechecks, [length t + mass t] does not. Values are\n stored in SI base units internally.\n\n Each dimension module provides three families of functions:\n - {e Scalar constructors} ([Length.km], [Mass.kg], ...) create a 0-d\n quantity from a [float].\n - {e Tensor constructors} ([Length.of_km], [Mass.of_kg], ...) wrap an\n arbitrary-shape {!Nx.float64_t}.\n - {e Extractors} ([Length.in_km], [Mass.in_kg], ...) return the numeric\n value in a given unit as an {!Nx.float64_t}.\n\n {[\n open Unit\n\n let d = Length.(kpc 10.0 + pc 500.0)\n let d_mpc = Length.in_mpc d\n ]} *)\n\n(** {1:types Types} *)\n\ntype 'a t\n(** The type for a physical quantity with dimension ['a]. Internally an\n {!Nx.float64_t} in SI base units. *)\n\ntype length\n(** Phantom type for length (SI: metres). *)\n\ntype mass\n(** Phantom type for mass (SI: kilograms). *)\n\ntype time\n(** Phantom type for time duration (SI: seconds). *)\n\ntype angle\n(** Phantom type for angles (SI: radians). *)\n\ntype velocity\n(** Phantom type for velocity (SI: m/s). *)\n\ntype power\n(** Phantom type for power / luminosity (SI: watts). *)\n\ntype temperature\n(** Phantom type for temperature (SI: kelvin). *)\n\ntype energy\n(** Phantom type for energy (SI: joules). *)\n\ntype frequency\n(** Phantom type for frequency (SI: hertz). *)\n\ntype dimensionless\n(** Phantom type for dimensionless quantities. *)\n\n(** {1:arithmetic Arithmetic}\n\n All operations require matching dimensions. *)\n\nval ( + ) : 'a t -> 'a t -> 'a t\n(** [a + b] is the element-wise sum of [a] and [b]. *)\n\nval ( - ) : 'a t -> 'a t -> 'a t\n(** [a - b] is the element-wise difference of [a] and [b]. *)\n\nval neg : 'a t -> 'a t\n(** [neg x] is the element-wise negation of [x]. *)\n\nval abs : 'a t -> 'a t\n(** [abs x] is the element-wise absolute value of [x]. *)\n\nval scale : float -> 'a t -> 'a t\n(** [scale s x] multiplies every element of [x] by [s]. *)\n\nval scale_t : Nx.float64_t -> 'a t -> 'a t\n(** [scale_t s x] multiplies every element of [x] by the tensor [s]. Keeps the\n result in the typed world when the scale factor is a fitted parameter. *)\n\nval ratio : 'a t -> 'a t -> dimensionless t\n(** [ratio a b] is the element-wise division [a / b], yielding a dimensionless\n quantity. *)\n\nval zero : 'a t\n(** [zero] is the scalar quantity [0.0]. *)\n\n(** {1:predicates Predicates, comparisons, and converting}\n\n These functions extract scalar values and are intended for 0-d tensors. *)\n\nval compare : 'a t -> 'a t -> int\n(** [compare a b] orders [a] and [b] by their scalar SI values. *)\n\nval equal : 'a t -> 'a t -> bool\n(** [equal a b] is [true] iff [a] and [b] have the same scalar SI value. *)\n\nval pp : Format.formatter -> 'a t -> unit\n(** [pp] formats the scalar SI value of a quantity. *)\n\nval to_float : 'a t -> float\n(** [to_float x] is the scalar value of [x] in SI base units. *)\n\n(** {1:cross Cross-dimension combinators}\n\n Functions that relate quantities of different dimensions. *)\n\nval length_per_time : length t -> time t -> velocity t\n(** [length_per_time d t] is [d / t] as a velocity. *)\n\nval velocity_times_time : velocity t -> time t -> length t\n(** [velocity_times_time v t] is [v * t] as a length. *)\n\nval length_per_velocity : length t -> velocity t -> time t\n(** [length_per_velocity d v] is [d / v] as a time. *)\n\nval wavelength_to_frequency : length t -> frequency t\n(** [wavelength_to_frequency lam] is [c / lam]. *)\n\nval frequency_to_wavelength : frequency t -> length t\n(** [frequency_to_wavelength nu] is [c / nu]. *)\n\nval frequency_to_energy : frequency t -> energy t\n(** [frequency_to_energy nu] is [h * nu]. *)\n\nval energy_to_frequency : energy t -> frequency t\n(** [energy_to_frequency e] is [e / h]. *)\n\nval energy_to_wavelength : energy t -> length t\n(** [energy_to_wavelength e] is [h * c / e]. *)\n\nval parallax_to_distance : angle t -> length t\n(** [parallax_to_distance p] is the distance corresponding to parallax [p]. Uses\n [d = 1 AU / p]. One arcsecond of parallax gives one parsec. *)\n\nval distance_to_parallax : length t -> angle t\n(** [distance_to_parallax d] is the parallax corresponding to distance [d]. Uses\n [p = 1 AU / d]. *)\n\nval flam_to_fnu : wavelength:length t -> Nx.float64_t -> Nx.float64_t\n(** [flam_to_fnu ~wavelength flam] converts spectral flux density from\n per-wavelength (F{_ {e lambda}}, W m{^ -2} m{^ -1}) to per-frequency\n (F{_ {e nu}}, W m{^ -2} Hz{^ -1}) at the given wavelengths. Uses\n [f_nu = f_lambda * lambda{^ 2} / c]. *)\n\nval fnu_to_flam : wavelength:length t -> Nx.float64_t -> Nx.float64_t\n(** [fnu_to_flam ~wavelength fnu] converts spectral flux density from\n per-frequency (F{_ {e nu}}) to per-wavelength (F{_ {e lambda}}) at the given\n wavelengths. Uses [f_lambda = f_nu * c / lambda{^ 2}]. *)\n\n(** {2:doppler Doppler conventions}\n\n Three conventions for converting between radial velocity and observed\n wavelength, given a rest wavelength. All agree at [v << c]. *)\n\nval doppler_optical : rest:length t -> velocity t -> length t\n(** [doppler_optical ~rest v] is the observed wavelength under the optical (cz)\n convention: [lambda_obs = lambda_rest * (1 + v/c)]. *)\n\nval doppler_optical_inv : rest:length t -> length t -> velocity t\n(** [doppler_optical_inv ~rest obs] is the velocity under the optical\n convention: [v = c * (lambda_obs/lambda_rest - 1)]. *)\n\nval doppler_radio : rest:length t -> velocity t -> length t\n(** [doppler_radio ~rest v] is the observed wavelength under the radio\n convention: [lambda_obs = lambda_rest / (1 - v/c)]. *)\n\nval doppler_radio_inv : rest:length t -> length t -> velocity t\n(** [doppler_radio_inv ~rest obs] is the velocity under the radio convention:\n [v = c * (1 - lambda_rest/lambda_obs)]. *)\n\nval doppler_relativistic : rest:length t -> velocity t -> length t\n(** [doppler_relativistic ~rest v] is the observed wavelength under the full\n relativistic Doppler formula:\n [lambda_obs = lambda_rest * sqrt((1 + v/c) / (1 - v/c))]. *)\n\nval doppler_relativistic_inv : rest:length t -> length t -> velocity t\n(** [doppler_relativistic_inv ~rest obs] is the velocity under the relativistic\n formula: [v = c * (r{^ 2} - 1) / (r{^ 2} + 1)] where\n [r = lambda_obs/lambda_rest]. *)\n\n(** {1:length Length} *)\n\nmodule Length : sig\n val of_tensor : Nx.float64_t -> length t\n (** [of_tensor x] wraps [x] as a length. [x] must be in metres. *)\n\n val to_tensor : length t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in metres. *)\n\n (** {2:scalar Scalar constructors}\n\n Each function creates a 0-d length quantity from a [float] value in the\n named unit. *)\n\n val m : float -> length t\n (** [m x] is [x] metres. *)\n\n val km : float -> length t\n (** [km x] is [x] kilometres. *)\n\n val cm : float -> length t\n (** [cm x] is [x] centimetres. *)\n\n val mm : float -> length t\n (** [mm x] is [x] millimetres. *)\n\n val um : float -> length t\n (** [um x] is [x] micrometres. *)\n\n val nm : float -> length t\n (** [nm x] is [x] nanometres. *)\n\n val angstrom : float -> length t\n (** [angstrom x] is [x] angstroms. *)\n\n val au : float -> length t\n (** [au x] is [x] astronomical units. *)\n\n val pc : float -> length t\n (** [pc x] is [x] parsecs. *)\n\n val kpc : float -> length t\n (** [kpc x] is [x] kiloparsecs. *)\n\n val mpc : float -> length t\n (** [mpc x] is [x] megaparsecs. *)\n\n val gpc : float -> length t\n (** [gpc x] is [x] gigaparsecs. *)\n\n val ly : float -> length t\n (** [ly x] is [x] light-years. *)\n\n val solar_radius : float -> length t\n (** [solar_radius x] is [x] solar radii. *)\n\n val earth_radius : float -> length t\n (** [earth_radius x] is [x] Earth equatorial radii. *)\n\n val jupiter_radius : float -> length t\n (** [jupiter_radius x] is [x] Jupiter equatorial radii. *)\n\n (** {2:tensor Tensor constructors}\n\n Each function wraps an arbitrary-shape {!Nx.float64_t} (in the named unit)\n as a length quantity. *)\n\n val of_m : Nx.float64_t -> length t\n val of_km : Nx.float64_t -> length t\n val of_cm : Nx.float64_t -> length t\n val of_mm : Nx.float64_t -> length t\n val of_um : Nx.float64_t -> length t\n val of_nm : Nx.float64_t -> length t\n val of_angstrom : Nx.float64_t -> length t\n val of_au : Nx.float64_t -> length t\n val of_pc : Nx.float64_t -> length t\n val of_kpc : Nx.float64_t -> length t\n val of_mpc : Nx.float64_t -> length t\n val of_gpc : Nx.float64_t -> length t\n val of_ly : Nx.float64_t -> length t\n val of_solar_radius : Nx.float64_t -> length t\n val of_earth_radius : Nx.float64_t -> length t\n val of_jupiter_radius : Nx.float64_t -> length t\n\n (** {2:extract Extracting}\n\n Each function returns the numeric value in the named unit as an\n {!Nx.float64_t}. *)\n\n val in_m : length t -> Nx.float64_t\n val in_km : length t -> Nx.float64_t\n val in_cm : length t -> Nx.float64_t\n val in_mm : length t -> Nx.float64_t\n val in_um : length t -> Nx.float64_t\n val in_nm : length t -> Nx.float64_t\n val in_angstrom : length t -> Nx.float64_t\n val in_au : length t -> Nx.float64_t\n val in_pc : length t -> Nx.float64_t\n val in_kpc : length t -> Nx.float64_t\n val in_mpc : length t -> Nx.float64_t\n val in_gpc : length t -> Nx.float64_t\n val in_ly : length t -> Nx.float64_t\n val in_solar_radius : length t -> Nx.float64_t\n val in_earth_radius : length t -> Nx.float64_t\n val in_jupiter_radius : length t -> Nx.float64_t\nend\n\n(** {1:mass Mass} *)\n\nmodule Mass : sig\n val of_tensor : Nx.float64_t -> mass t\n (** [of_tensor x] wraps [x] as a mass. [x] must be in kilograms. *)\n\n val to_tensor : mass t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in kilograms. *)\n\n (** {2:scalar Scalar constructors} *)\n\n val kg : float -> mass t\n (** [kg x] is [x] kilograms. *)\n\n val g : float -> mass t\n (** [g x] is [x] grams. *)\n\n val mg : float -> mass t\n (** [mg x] is [x] milligrams. *)\n\n val solar_mass : float -> mass t\n (** [solar_mass x] is [x] solar masses. *)\n\n val earth_mass : float -> mass t\n (** [earth_mass x] is [x] Earth masses. *)\n\n val jupiter_mass : float -> mass t\n (** [jupiter_mass x] is [x] Jupiter masses. *)\n\n (** {2:tensor Tensor constructors} *)\n\n val of_kg : Nx.float64_t -> mass t\n val of_g : Nx.float64_t -> mass t\n val of_mg : Nx.float64_t -> mass t\n val of_solar_mass : Nx.float64_t -> mass t\n val of_earth_mass : Nx.float64_t -> mass t\n val of_jupiter_mass : Nx.float64_t -> mass t\n\n (** {2:extract Extracting} *)\n\n val in_kg : mass t -> Nx.float64_t\n val in_g : mass t -> Nx.float64_t\n val in_mg : mass t -> Nx.float64_t\n val in_solar_mass : mass t -> Nx.float64_t\n val in_earth_mass : mass t -> Nx.float64_t\n val in_jupiter_mass : mass t -> Nx.float64_t\nend\n\n(** {1:time Time duration} *)\n\nmodule Time : sig\n val of_tensor : Nx.float64_t -> time t\n (** [of_tensor x] wraps [x] as a time duration. [x] must be in seconds. *)\n\n val to_tensor : time t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in seconds. *)\n\n (** {2:scalar Scalar constructors} *)\n\n val s : float -> time t\n (** [s x] is [x] seconds. *)\n\n val ms : float -> time t\n (** [ms x] is [x] milliseconds. *)\n\n val us : float -> time t\n (** [us x] is [x] microseconds. *)\n\n val min : float -> time t\n (** [min x] is [x] minutes. *)\n\n val hr : float -> time t\n (** [hr x] is [x] hours. *)\n\n val day : float -> time t\n (** [day x] is [x] days (86 400 s). *)\n\n val yr : float -> time t\n (** [yr x] is [x] Julian years (365.25 days). *)\n\n val myr : float -> time t\n (** [myr x] is [x] megayears. *)\n\n val gyr : float -> time t\n (** [gyr x] is [x] gigayears. *)\n\n (** {2:tensor Tensor constructors} *)\n\n val of_s : Nx.float64_t -> time t\n val of_ms : Nx.float64_t -> time t\n val of_us : Nx.float64_t -> time t\n val of_min : Nx.float64_t -> time t\n val of_hr : Nx.float64_t -> time t\n val of_day : Nx.float64_t -> time t\n val of_yr : Nx.float64_t -> time t\n val of_myr : Nx.float64_t -> time t\n val of_gyr : Nx.float64_t -> time t\n\n (** {2:extract Extracting} *)\n\n val in_s : time t -> Nx.float64_t\n val in_ms : time t -> Nx.float64_t\n val in_us : time t -> Nx.float64_t\n val in_min : time t -> Nx.float64_t\n val in_hr : time t -> Nx.float64_t\n val in_day : time t -> Nx.float64_t\n val in_yr : time t -> Nx.float64_t\n val in_myr : time t -> Nx.float64_t\n val in_gyr : time t -> Nx.float64_t\nend\n\n(** {1:angle Angle} *)\n\nmodule Angle : sig\n val of_tensor : Nx.float64_t -> angle t\n (** [of_tensor x] wraps [x] as an angle. [x] must be in radians. *)\n\n val to_tensor : angle t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in radians. *)\n\n (** {2:scalar Scalar constructors} *)\n\n val rad : float -> angle t\n (** [rad x] is [x] radians. *)\n\n val deg : float -> angle t\n (** [deg x] is [x] degrees. *)\n\n val arcmin : float -> angle t\n (** [arcmin x] is [x] arcminutes. *)\n\n val arcsec : float -> angle t\n (** [arcsec x] is [x] arcseconds. *)\n\n val mas : float -> angle t\n (** [mas x] is [x] milliarcseconds. *)\n\n val hour_angle : float -> angle t\n (** [hour_angle x] is [x] hour angles (1 h = 15 deg). *)\n\n (** {2:tensor Tensor constructors} *)\n\n val of_rad : Nx.float64_t -> angle t\n val of_deg : Nx.float64_t -> angle t\n val of_arcmin : Nx.float64_t -> angle t\n val of_arcsec : Nx.float64_t -> angle t\n val of_mas : Nx.float64_t -> angle t\n val of_hour_angle : Nx.float64_t -> angle t\n\n (** {2:extract Extracting} *)\n\n val in_rad : angle t -> Nx.float64_t\n val in_deg : angle t -> Nx.float64_t\n val in_arcmin : angle t -> Nx.float64_t\n val in_arcsec : angle t -> Nx.float64_t\n val in_mas : angle t -> Nx.float64_t\n val in_hour_angle : angle t -> Nx.float64_t\n\n (** {2:trig Trigonometric functions} *)\n\n val sin : angle t -> Nx.float64_t\n (** [sin a] is the sine of [a]. *)\n\n val cos : angle t -> Nx.float64_t\n (** [cos a] is the cosine of [a]. *)\n\n val tan : angle t -> Nx.float64_t\n (** [tan a] is the tangent of [a]. *)\n\n val asin : Nx.float64_t -> angle t\n (** [asin x] is the arc sine of [x]. *)\n\n val acos : Nx.float64_t -> angle t\n (** [acos x] is the arc cosine of [x]. *)\n\n val atan2 : y:Nx.float64_t -> x:Nx.float64_t -> angle t\n (** [atan2 ~y ~x] is the two-argument arc tangent of [y] and [x]. *)\n\n (** {2:wrap Wrapping} *)\n\n val wrap_360 : angle t -> angle t\n (** [wrap_360 a] normalizes [a] into \\[0, 360) degrees. *)\n\n val wrap_180 : angle t -> angle t\n (** [wrap_180 a] normalizes [a] into \\[-180, 180) degrees. *)\nend\n\n(** {1:velocity Velocity} *)\n\nmodule Velocity : sig\n val of_tensor : Nx.float64_t -> velocity t\n (** [of_tensor x] wraps [x] as a velocity. [x] must be in m/s. *)\n\n val to_tensor : velocity t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in m/s. *)\n\n (** {2:scalar Scalar constructors} *)\n\n val m_s : float -> velocity t\n (** [m_s x] is [x] m/s. *)\n\n val km_s : float -> velocity t\n (** [km_s x] is [x] km/s. *)\n\n val km_hr : float -> velocity t\n (** [km_hr x] is [x] km/h. *)\n\n (** {2:tensor Tensor constructors} *)\n\n val of_m_s : Nx.float64_t -> velocity t\n val of_km_s : Nx.float64_t -> velocity t\n val of_km_hr : Nx.float64_t -> velocity t\n\n (** {2:extract Extracting} *)\n\n val in_m_s : velocity t -> Nx.float64_t\n val in_km_s : velocity t -> Nx.float64_t\n val in_km_hr : velocity t -> Nx.float64_t\nend\n\n(** {1:power Power / Luminosity} *)\n\nmodule Power : sig\n val of_tensor : Nx.float64_t -> power t\n (** [of_tensor x] wraps [x] as a power. [x] must be in watts. *)\n\n val to_tensor : power t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in watts. *)\n\n (** {2:scalar Scalar constructors} *)\n\n val w : float -> power t\n (** [w x] is [x] watts. *)\n\n val kw : float -> power t\n (** [kw x] is [x] kilowatts. *)\n\n val solar_luminosity : float -> power t\n (** [solar_luminosity x] is [x] solar luminosities. *)\n\n val erg_s : float -> power t\n (** [erg_s x] is [x] erg/s. *)\n\n (** {2:tensor Tensor constructors} *)\n\n val of_w : Nx.float64_t -> power t\n val of_kw : Nx.float64_t -> power t\n val of_solar_luminosity : Nx.float64_t -> power t\n val of_erg_s : Nx.float64_t -> power t\n\n (** {2:extract Extracting} *)\n\n val in_w : power t -> Nx.float64_t\n val in_kw : power t -> Nx.float64_t\n val in_solar_luminosity : power t -> Nx.float64_t\n val in_erg_s : power t -> Nx.float64_t\nend\n\n(** {1:temperature Temperature} *)\n\nmodule Temperature : sig\n val of_tensor : Nx.float64_t -> temperature t\n (** [of_tensor x] wraps [x] as a temperature. [x] must be in kelvin. *)\n\n val to_tensor : temperature t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in kelvin. *)\n\n (** {2:scalar Scalar constructors} *)\n\n val kelvin : float -> temperature t\n (** [kelvin x] is [x] kelvin. *)\n\n (** {2:tensor Tensor constructors} *)\n\n val of_kelvin : Nx.float64_t -> temperature t\n\n (** {2:extract Extracting} *)\n\n val in_kelvin : temperature t -> Nx.float64_t\nend\n\n(** {1:energy Energy} *)\n\nmodule Energy : sig\n val of_tensor : Nx.float64_t -> energy t\n (** [of_tensor x] wraps [x] as an energy. [x] must be in joules. *)\n\n val to_tensor : energy t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in joules. *)\n\n (** {2:scalar Scalar constructors} *)\n\n val j : float -> energy t\n (** [j x] is [x] joules. *)\n\n val erg : float -> energy t\n (** [erg x] is [x] ergs. *)\n\n val ev : float -> energy t\n (** [ev x] is [x] electronvolts. *)\n\n val kev : float -> energy t\n (** [kev x] is [x] kiloelectronvolts. *)\n\n val mev : float -> energy t\n (** [mev x] is [x] megaelectronvolts. *)\n\n (** {2:tensor Tensor constructors} *)\n\n val of_j : Nx.float64_t -> energy t\n val of_erg : Nx.float64_t -> energy t\n val of_ev : Nx.float64_t -> energy t\n val of_kev : Nx.float64_t -> energy t\n val of_mev : Nx.float64_t -> energy t\n\n (** {2:extract Extracting} *)\n\n val in_j : energy t -> Nx.float64_t\n val in_erg : energy t -> Nx.float64_t\n val in_ev : energy t -> Nx.float64_t\n val in_kev : energy t -> Nx.float64_t\n val in_mev : energy t -> Nx.float64_t\nend\n\n(** {1:frequency Frequency} *)\n\nmodule Frequency : sig\n val of_tensor : Nx.float64_t -> frequency t\n (** [of_tensor x] wraps [x] as a frequency. [x] must be in hertz. *)\n\n val to_tensor : frequency t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor in hertz. *)\n\n (** {2:scalar Scalar constructors} *)\n\n val hz : float -> frequency t\n (** [hz x] is [x] hertz. *)\n\n val khz : float -> frequency t\n (** [khz x] is [x] kilohertz. *)\n\n val mhz : float -> frequency t\n (** [mhz x] is [x] megahertz. *)\n\n val ghz : float -> frequency t\n (** [ghz x] is [x] gigahertz. *)\n\n (** {2:tensor Tensor constructors} *)\n\n val of_hz : Nx.float64_t -> frequency t\n val of_khz : Nx.float64_t -> frequency t\n val of_mhz : Nx.float64_t -> frequency t\n val of_ghz : Nx.float64_t -> frequency t\n\n (** {2:extract Extracting} *)\n\n val in_hz : frequency t -> Nx.float64_t\n val in_khz : frequency t -> Nx.float64_t\n val in_mhz : frequency t -> Nx.float64_t\n val in_ghz : frequency t -> Nx.float64_t\nend\n\n(** {1:dimensionless Dimensionless} *)\n\nmodule Dimensionless : sig\n val of_tensor : Nx.float64_t -> dimensionless t\n (** [of_tensor x] wraps [x] as a dimensionless quantity. *)\n\n val to_tensor : dimensionless t -> Nx.float64_t\n (** [to_tensor x] is the underlying tensor. *)\n\n val v : float -> dimensionless t\n (** [v x] is the scalar dimensionless quantity [x]. *)\n\n val to_float : dimensionless t -> float\n (** [to_float x] is the scalar value of [x]. Intended for 0-d tensors. *)\nend\n"
},
{
"path": "dev/umbra/lib/vega_data.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Vega (alpha Lyrae) reference spectrum from CALSPEC: alpha_lyr_stis_011.fits\n (Bohlin 2014). Subsampled to 300 points from 900-250000 Angstroms.\n\n wave: wavelength in Angstroms. flux: spectral flux density f_lambda in\n W/m^2/m (SI). *)\n\n[@@@ocamlformat \"disable\"]\n\nlet wave =\n [|\n 900.452; 917.759; 935.398; 953.377; 971.701;\n 989.386; 1008.402; 1027.784; 1047.539; 1067.673;\n 1087.104; 1107.998; 1129.294; 1151.0; 1172.199951171875;\n 1198.199951171875; 1230.0; 1240.699951171875; 1264.300048828125; 1287.9000244140625;\n 1312.5999755859375; 1337.4000244140625; 1363.4000244140625; 1388.0999755859375; 1415.300048828125;\n 1442.4000244140625; 1469.5; 1496.699951171875; 1526.199951171875; 1554.5;\n 1584.0; 1614.5999755859375; 1645.300048828125; 1676.115356; 1707.652832;\n 1740.569458; 1773.494019; 1807.798096; 1842.109619; 1876.428101;\n 1912.126099; 1947.830688; 1984.914551; 2023.37793; 2061.84668;\n 2100.320312; 2140.172119; 2181.401855; 2222.634521; 2265.243652;\n 2307.85376; 2351.838867; 2395.823242; 2441.179932; 2487.907959;\n 2536.005615; 2584.09668; 2632.18042; 2683.001953; 2733.811768;\n 2785.980957; 2838.134277; 2891.641846; 2946.500488; 3002.706055;\n 3060.254883; 3119.17334; 3176.848877; 3237.280029; 3300.467529;\n 3363.663574; 3426.86792; 3492.827881; 3558.794678; 3624.767822;\n 3693.495117; 3764.976562; 3836.461914; 3907.949951; 3982.189453;\n 4059.178955; 4136.16748; 4213.153809; 4292.885742; 4375.36084;\n 4457.828613; 4543.035156; 4630.977539; 4718.905273; 4806.816406;\n 4897.454102; 4990.81543; 5086.894531; 5182.942871; 5281.700684;\n 5380.419922; 5485.38916; 5587.699219; 5694.90625; 5802.137207;\n 5914.267578; 6026.421387; 6138.598145; 6255.674805; 6377.652344;\n 6494.770996; 6621.67041; 6743.70752; 6875.525879; 7002.478027;\n 7139.210449; 7271.072266; 7412.712402; 7549.477051; 7696.016602;\n 7842.55957; 7989.103027; 8140.52832; 8296.831055; 8453.124023;\n 8614.286133; 8775.428711; 8946.311523; 9112.282227; 9287.978516;\n 9463.631836; 9644.114258; 9824.539062; 10014.648438; 10205.378729443204;\n 10401.529770846228; 10601.450905715283; 10794.393204668508; 11001.865315614596; 11213.325115099247;\n 11428.849247658822; 11648.515831526549; 11860.514306414294; 12088.477647568725; 12320.82252599989;\n 12557.633156314389; 12786.177214329819; 13031.93212871661; 13282.410540749495; 13537.703237775104;\n 13797.902752017966; 14049.019236198084; 14319.046427392743; 14594.263638185937; 14874.770622496817;\n 15145.48568465375; 15436.587354440591; 15733.284102879634; 16035.6834692452; 16343.895059966944;\n 16641.34761761389; 16961.200290683963; 17287.200646834564; 17619.466846796713; 17940.134316746688;\n 18284.950136177056; 18636.393439411142; 18994.591609034607; 19359.674476061784; 19712.013050644106;\n 20090.8850151999; 20477.039034852016; 20870.61507330569; 21250.45218142002; 21658.893498090372;\n 22075.185203484638; 22499.47818484361; 22931.926229505145; 23349.278401987125; 23798.05991182052;\n 24255.46716331803; 24721.665946277586; 25171.590688206736; 25655.39769412964; 26148.503644305452;\n 26651.08726772658; 27163.33072884503; 27657.69282371489; 28189.283604893368; 28731.09175155295;\n 29283.313645250695; 29816.258607862546; 30389.33779785106; 30973.431775616467; 31568.752249243276;\n 32175.51499603123; 32761.096902670255; 33390.77694148892; 34032.559656653226; 34686.6776658844;\n 35317.961760150065; 35996.78565914454; 36688.65679708345; 37393.82594651755; 38112.548700060266;\n 38806.182319759384; 39552.05106970492; 40312.25568451587; 41087.071705009694; 41834.84049088929;\n 42638.92113772262; 43458.45650302316; 44293.74363220345; 45145.08527967651; 45966.708340478566;\n 46850.20497002188; 47750.6827218135; 48668.46797921292; 49603.893398540575; 50506.66503950412;\n 51477.42126605539; 52466.83577548603; 53475.2671869996; 54448.49633089426; 55495.01596936514;\n 56561.650090809424; 57648.78530287201; 58756.81564377781; 59826.16692337332; 60976.04732038082;\n 62148.02883123579; 63342.53624731101; 64495.34586512943; 65734.96955693385; 66998.41926097606;\n 68286.15292170878; 69598.63728556872; 70865.30551837588; 72227.3621205334; 73615.59793944974;\n 75030.51614901233; 76396.04247255616; 77864.4018237676; 79360.98356865476; 80886.33015165699;\n 82440.99444348396; 83941.38860077565; 85554.77221228057; 87199.16563569565; 88875.16489081086;\n 90492.65845739808; 92231.95982648393; 94004.69119201355; 95811.49509052176; 97653.02640827911;\n 99430.27365358408; 101341.3591939848; 103289.17648427008; 105274.43152183376; 107190.38538724542;\n 109250.62305986807; 111350.45923995743; 113490.65502485144; 115671.98614086409; 117777.17147181589;\n 120040.89097212761; 122348.11997351186; 124699.69474434588; 126969.18087410457; 129409.57409375694;\n 131896.8725467459; 134431.9777685813; 137015.80862256046; 139509.44325303897; 142190.86481564713;\n 144923.8242628955; 147709.31217146275; 150397.56442345414; 153288.25958223912; 156234.51493917976;\n 159237.39838144454; 162297.9983210655; 165251.7590085775; 168427.95711765194; 171665.20289426477;\n 174964.66969630358; 178148.95889525744; 181573.045814387; 185062.94491283095; 188619.92112426818;\n 192245.26369497634; 195744.05300937625; 199506.32395143431; 203340.90709103053; 207249.19229470106;\n 211021.04729625938; 215076.94755071797; 219210.80366346712; 223424.11397425184; 227718.40562110202;\n 231862.792374509; 236319.27844847148; 240861.41978391458; 245490.86270582714; 249958.7012923323\n |]\n\nlet flux =\n [|\n 1.2837948759614193e-10; 1.1943805588998657e-07; 1.173052282865683e-06; 2.1574317088379757e-06; 1.3420374216366326e-06;\n 4.100760634173639e-06; 6.2335830079973675e-06; 4.519122285273625e-06; 1.093481569114374e-05; 1.3420373761618976e-05;\n 1.8457114492775872e-05; 4.486309626372531e-05; 0.0006154832663014531; 0.0017448125872761011; 0.00016285567835438997;\n 0.00012107902148272842; 0.0002881045511458069; 0.0008428706787526608; 0.00555039057508111; 0.021335331723093987;\n 0.032872095704078674; 0.04942339286208153; 0.05074312165379524; 0.05717964470386505; 0.062469806522130966;\n 0.0662260353565216; 0.06470223516225815; 0.0783146470785141; 0.06845968961715698; 0.07111997902393341;\n 0.0734374150633812; 0.0725170373916626; 0.07288476824760437; 0.06900300085544586; 0.0641390010714531;\n 0.06533699482679367; 0.06514099985361099; 0.06312000006437302; 0.058240000158548355; 0.060812000185251236;\n 0.059772998094558716; 0.060645997524261475; 0.05867400020360947; 0.0558370016515255; 0.05407499894499779;\n 0.05430099740624428; 0.04962899908423424; 0.05051000043749809; 0.047245997935533524; 0.044964998960494995;\n 0.04547500237822533; 0.04311799630522728; 0.03780499845743179; 0.041788000613451004; 0.04115099832415581;\n 0.03668300062417984; 0.03907399997115135; 0.03461499884724617; 0.040741000324487686; 0.03929299861192703;\n 0.03749600052833557; 0.037234000861644745; 0.03612099960446358; 0.03686100244522095; 0.03653800114989281;\n 0.036056000739336014; 0.03495499864220619; 0.03484500199556351; 0.033263999968767166; 0.03422100096940994;\n 0.032947998493909836; 0.032896000891923904; 0.03175799921154976; 0.031553998589515686; 0.03090999834239483;\n 0.0309200007468462; 0.045221999287605286; 0.033542998135089874; 0.07968499511480331; 0.06394299864768982;\n 0.08404000103473663; 0.07898300141096115; 0.07822800427675247; 0.0719740018248558; 0.0667010024189949;\n 0.06643600016832352; 0.06298799812793732; 0.05953500047326088; 0.05647499859333038; 0.05225900188088417;\n 0.047912996262311935; 0.04777200147509575; 0.04543299973011017; 0.04206399992108345; 0.04022299870848656;\n 0.03820599988102913; 0.03601999953389168; 0.03411700204014778; 0.032175999134778976; 0.03042599931359291;\n 0.028788000345230103; 0.027058999985456467; 0.02562600001692772; 0.024166999384760857; 0.02274700067937374;\n 0.021289000287652016; 0.019968999549746513; 0.019201001152396202; 0.01802999898791313; 0.017007999122142792;\n 0.01603100076317787; 0.015209999866783619; 0.014260999858379364; 0.013499000109732151; 0.012649999931454659;\n 0.011848000809550285; 0.011309999972581863; 0.010559000074863434; 0.009970699436962605; 0.009328600019216537;\n 0.0090791005641222; 0.008882399648427963; 0.009110400453209877; 0.008642599917948246; 0.008087899535894394;\n 0.007705099880695343; 0.00721339974552393; 0.006843299604952335; 0.006151599809527397; 0.006024217698723078;\n 0.005648050922900438; 0.0052790273912250996; 0.004945714958012104; 0.004534630570560694; 0.0043441662564873695;\n 0.004064818844199181; 0.0037846784107387066; 0.0035608832258731127; 0.0033259775955229998; 0.0031156735494732857;\n 0.0029069569427520037; 0.002545868745073676; 0.002553011057898402; 0.0023879422806203365; 0.0022300160489976406;\n 0.0020839935168623924; 0.001953843282535672; 0.0018244864186272025; 0.0017038591904565692; 0.0015895807882770896;\n 0.0014951423509046435; 0.0013784831389784813; 0.0012856320245191455; 0.0012396031524986029; 0.0011531008640304208;\n 0.0011078655952587724; 0.0010332672391086817; 0.0009435904212296009; 0.0008983552106656134; 0.0008412160095758736;\n 0.0007756645791232586; 0.0007241599960252643; 0.0006796390516683459; 0.0006226585246622562; 0.0005917081143707037;\n 0.0005520281265489757; 0.0005140146822668612; 0.0004784614429809153; 0.00044632062781602144; 0.00036132606328465044;\n 0.0003871974186040461; 0.0003606118552852422; 0.0003356928064022213; 0.0003139481705147773; 0.00029244160396046937;\n 0.0002722047793213278; 0.0002522854192648083; 0.0002364928077440709; 0.00021379583631642163; 0.00019514624727889895;\n 0.00018959103908855468; 0.00017943295824807137; 0.0001626086450414732; 0.00015602176426909864; 0.00013110271538607776;\n 0.0001348326331935823; 0.0001258649572264403; 0.0001058662383002229; 0.00010872319398913532; 0.00010110463335877284;\n 9.396224049851298e-05; 8.72959935804829e-05; 8.126463944790885e-05; 7.560627273051068e-05; 7.031296263448894e-05;\n 6.562278576893732e-05; 6.097228470025584e-05; 5.6718592531979084e-05; 4.850483310292475e-05; 4.869529584539123e-05;\n 4.574310514726676e-05; 4.2584575567161664e-05; 3.8981630495982245e-05; 3.6783359973924235e-05; 3.387084507266991e-05;\n 3.1902720365906134e-05; 2.955366471724119e-05; 2.7506175683811307e-05; 2.5553918021614663e-05; 2.3807999241398647e-05;\n 2.1720832592109218e-05; 2.0554240109049715e-05; 1.9094017261522822e-05; 1.7729023966239765e-05; 1.6522752048331313e-05;\n 1.520537625765428e-05; 1.422924833605066e-05; 1.3229311662144028e-05; 1.2340479770500679e-05; 1.143577628681669e-05;\n 1.063423951563891e-05; 9.872383998299483e-06; 9.134336323768366e-06; 8.507391612511128e-06; 7.931238542369101e-06;\n 7.359846222243505e-06; 6.820991984568536e-06; 6.3662591855973005e-06; 5.908352250116877e-06; 5.479807896335842e-06;\n 5.0814210226235446e-06; 4.719538992503658e-06; 4.393369636090938e-06; 4.080691269336967e-06; 3.7862655517528765e-06;\n 3.291852635811665e-06; 3.273599986641784e-06; 3.0180608519003727e-06; 2.8196607217978453e-06; 2.615705398056889e-06;\n 2.4268288143503014e-06; 2.2609663119510515e-06; 2.0966911051800707e-06; 1.937177557920222e-06; 1.8046463310383842e-06;\n 1.6800512412373791e-06; 1.5578368675051024e-06; 1.4118143099040026e-06; 1.3411839745458565e-06; 1.2443647392501589e-06;\n 1.1586560049181571e-06; 1.0753279866548837e-06; 9.967616279027425e-07; 9.150207347374817e-07; 8.610560371380416e-07;\n 7.983616114870529e-07; 7.404287885037775e-07; 6.840038508926227e-07; 6.342451115415315e-07; 5.927398092353542e-07;\n 5.49964795482083e-07; 5.098086148791481e-07; 4.7227135269167775e-07; 4.399718420700083e-07; 4.0822783375915606e-07;\n 3.7775359373881656e-07; 3.5092992334284645e-07; 3.2537599281567964e-07; 3.019648033841804e-07; 2.7887102760359994e-07;\n 2.6038017608698283e-07; 2.414131188288593e-07; 2.2458880266640335e-07; 2.0831998881476466e-07; 1.930828688045949e-07;\n 1.790361494613535e-07; 1.6260864299511013e-07; 1.5435520595019625e-07; 1.4181631513565662e-07; 1.32769287120027e-07;\n 1.230079931247019e-07; 1.1451648163074424e-07; 1.0618367696224595e-07; 9.83270425081173e-08; 9.118463850654734e-08;\n 8.451839761391966e-08; 7.86695650845104e-08; 7.293977688505038e-08; 6.760678417094823e-08; 6.264678376055599e-08;\n 5.8289920445986354e-08; 5.4044157593580167e-08; 5.0052349820361997e-08; 4.447334234214395e-08; 4.3013120176738084e-08;\n 4.002918529977251e-08; 3.6950016379933004e-08; 3.4370817303397416e-08; 3.1863038429946755e-08; 2.9625088160400992e-08\n |]\n"
},
{
"path": "dev/umbra/papers/perlmutter1999/.gitignore",
"content": "/data/\n"
},
{
"path": "dev/umbra/papers/perlmutter1999/download_data.sh",
"content": "#!/usr/bin/env bash\n# Download Pantheon+ Type Ia supernova data (Scolnic et al. 2022, Brout et al. 2022).\n#\n# The Pantheon+ compilation contains 1701 light curves of 1550 unique SNe Ia\n# spanning 0.001 < z < 2.26, extending the original 42 high-z supernovae from\n# Perlmutter et al. (1999) that first demonstrated cosmic acceleration.\n#\n# Source: https://github.com/PantheonPlusSH0ES/DataRelease\n# Papers: arXiv:2112.03863 (data), arXiv:2202.04077 (cosmology)\n\nset -euo pipefail\n\nDIR=\"$(cd \"$(dirname \"$0\")\" && pwd)\"\nDATA_DIR=\"${DIR}/data\"\nmkdir -p \"${DATA_DIR}\"\n\nBASE_URL=\"https://raw.githubusercontent.com/PantheonPlusSH0ES/DataRelease/main/Pantheon%2B_Data/4_DISTANCES_AND_COVAR\"\n\necho \"Downloading Pantheon+ SN Ia distance data...\"\ncurl -fSL \"${BASE_URL}/Pantheon%2BSH0ES.dat\" -o \"${DATA_DIR}/Pantheon+SH0ES.dat\"\necho \" -> ${DATA_DIR}/Pantheon+SH0ES.dat ($(wc -l < \"${DATA_DIR}/Pantheon+SH0ES.dat\") lines)\"\n\necho \"Downloading paper PDF (Perlmutter et al. 1999, arXiv:astro-ph/9812133)...\"\ncurl -fSL \"https://arxiv.org/pdf/astro-ph/9812133\" -o \"${DATA_DIR}/perlmutter1999.pdf\"\necho \" -> ${DATA_DIR}/perlmutter1999.pdf\"\n\necho \"Done.\"\n"
},
{
"path": "dev/umbra/papers/perlmutter1999/perlmutter1999.md",
"content": "# The Accelerating Universe\n\nReproducing the key result of Perlmutter et al. (1999), \"Measurements of\n$\\Omega$ and $\\Lambda$ from 42 High-Redshift Supernovae\" (ApJ 517, 565) --\nthe Nobel Prize-winning discovery that the expansion of the universe is\naccelerating.\n\nWe use the modern Pantheon+ dataset (Scolnic et al. 2022, 1701 SNe Ia spanning\n$0.001 < z < 2.26$) which extends the original 42 supernovae and confirms the\nresult with far greater precision.\n\n\n## Background\n\nType Ia supernovae (SNe Ia) are \"standardizable candles\": after correcting for\nthe correlation between peak luminosity and light-curve width, they have\nremarkably uniform absolute magnitudes. This lets us measure their distances\nthrough the **distance modulus**:\n\n$$\\mu = m - M = 5 \\log_{10}\\!\\left(\\frac{d_L}{\\text{Mpc}}\\right) + 25$$\n\nwhere $d_L$ is the luminosity distance, which depends on the cosmological\nparameters $\\Omega_M$ (matter density) and $\\Omega_\\Lambda$ (dark energy\ndensity). In 1998--1999, two independent teams (the Supernova Cosmology Project\nand the High-z Supernova Search Team) found that distant SNe Ia are **fainter\nthan expected** in a decelerating universe -- implying that the expansion is\naccelerating, driven by a cosmological constant or dark energy.\n\nWe reproduce three key results:\n1. The **Hubble diagram** ($\\mu$ vs $z$) with cosmological model curves\n2. **Residuals** relative to an empty universe, showing the acceleration signal\n3. **Confidence contours** in the $\\Omega_M$--$\\Omega_\\Lambda$ plane\n\n\n## Setup\n\n\n\n```ocaml\n#require \"umbra\";;\n\nopen Nx\nopen Umbra\n\nlet f64 = Nx.float64\nlet f32 = Nx.float32\n```\n\n\nval f64 : (float, Nx.float64_elt) Nx.dtype = Nx.Float64\nval f32 : (float, Nx.float32_elt) Nx.dtype = Nx.Float32\n\n\n## Loading the Pantheon+ data\n\nThe Pantheon+ compilation (Scolnic et al. 2022) provides standardized distance\nmoduli for 1701 SN Ia light curves from 18 surveys. We load the data file\n(downloaded by `download_data.sh`) and extract redshift, distance modulus, and\nthe diagonal error (for plotting; full cosmological fits require the covariance\nmatrix).\n\n\n\n```ocaml\nlet df = Talon_csv.read ~sep:' ' \"data/Pantheon+SH0ES.dat\"\n\nlet () =\n Printf.printf \"Loaded %d light curves, %d columns\\n\"\n (Talon.num_rows df) (List.length (Talon.column_names df))\n```\n\n\nLoaded 1701 light curves, 47 columns\nval df : Talon.t =\n \n\n\nCID IDSURVEY zHD zHDERR zCMB zCMBERR zHEL zHELERR m_b_corr m_b_corr_err_DIAG … \n2011fe 51 0.00122 0.00084 0.00122 2e-05 0.00082 2e-05 9.74571 1.51621 … 2011fe 56 0.00122 0.00084 0.00122 2e-05 0.00082 2e-05 9.80286 1.51723 … 2012cg 51 0.00256 0.00084 0.00256 2e-05 0.00144 2e-05 11.4703 0.781906 … 2012cg 56 0.00256 0.00084 0.00256 2e-05 0.00144 2e-05 11.4919 0.798612 … 1994DRichmond 50 0.00299 0.00084 0.00299 4e-05 0.00187 4e-05 11.5227 0.880798 … 1981B 50 0.00317 0.00084 0.0035 1e-05 0.00236 1e-05 11.5416 0.613941 … 2013aa 56 0.00331 0.00085 0.00478 0.00015 0.00411 0.00015 11.2074 0.59407 … 2013aa 5 0.00331 0.00085 0.00478 0.00015 0.00411 0.00015 11.2998 0.579622 … 2017cbv 5 0.00331 0.00085 0.00478 0.00015 0.00411 0.00015 11.1483 0.577815 … 2017cbv 18 0.00331 0.00085 0.00478 0.00015 0.00411 0.00015 11.2577 0.577916 … 2001el 50 0.00333 0.00084 0.00357 1e-05 0.00379 1e-05 12.2481 0.590389 … 2011by 51 0.00349 0.00084 0.00369 2e-05 0.00313 2e-05 12.5403 0.55206 … 1998aq 50 0.00349 0.00084 0.00369 1e-05 0.00313 1e-05 12.2437 0.544824 … 1990N 50 0.00359 0.00084 0.00462 2e-05 0.00355 2e-05 12.4439 0.550332 … 2021pit 56 0.00384 0.00084 0.00366 1e-05 0.00388 1e-05 11.7469 0.565861 … 2005df 50 0.00407 0.00084 0.00435 1e-05 0.00435 1e-05 12.1403 0.475638 … 2005df_ANU 50 0.00407 0.00084 0.00435 1e-05 0.00435 1e-05 12.1249 0.478515 … 2013dy 51 0.00432 0.00084 0.00293 0.00012 0.00394 0.00012 12.246 0.513549 … 2013dy 56 0.00432 0.00084 0.00293 0.00012 0.00394 0.00012 12.3081 0.530151 … 2012ht 56 0.00465 0.00084 0.00465 2e-05 0.00352 2e-05 12.6779 0.441191 … … … … … … … … … … … … \n
\n1701 rows × 47 columns
\n\n\n\n```ocaml\nlet col name =\n Talon.get_column_exn df name |> Talon.Col.to_tensor f64 |> Option.get\n\nlet sn_z = col \"zHD\"\nlet sn_mu = col \"MU_SH0ES\"\nlet sn_mu_err = col \"MU_SH0ES_ERR_DIAG\"\n\nlet () =\n let n = (Nx.shape sn_z).(0) in\n Printf.printf \"%d SNe Ia, z in [%.4f, %.3f]\\n\" n\n (Nx.item [] (Nx.min sn_z)) (Nx.item [] (Nx.max sn_z))\n```\n\n\n1701 SNe Ia, z in [0.0012, 2.261]\nval col : string -> (float, Nx.float64_elt) Nx.t = \nval sn_z : (float, Nx.float64_elt) Nx.t = float64 [1701] \n [0.00122, 0.00122, ..., 1.91165, 2.26137]\nval sn_mu : (float, Nx.float64_elt) Nx.t = float64 [1701] \n [28.9987, 29.0559, ..., 45.4233, 46.1828]\nval sn_mu_err : (float, Nx.float64_elt) Nx.t = float64 [1701] \n [1.51645, 1.51747, ..., 0.358642, 0.281309]\n\n\n## Cosmological models\n\nWe compute the theoretical distance modulus $\\mu(z)$ for several cosmologies\nto compare with the data. The key insight from Perlmutter et al. is that the\ndata prefer $\\Omega_\\Lambda > 0$ (accelerating expansion) over\n$\\Omega_\\Lambda = 0$ (decelerating expansion).\n\nThe models we compare:\n- **Best-fit $\\Lambda$CDM**: $(\\Omega_M, \\Omega_\\Lambda) = (0.3, 0.7)$, $H_0 = 70$\n- **Einstein--de Sitter**: $(\\Omega_M, \\Omega_\\Lambda) = (1, 0)$ -- matter-only, decelerating\n- **Empty (Milne)**: $(\\Omega_M, \\Omega_\\Lambda) = (0, 0)$ -- coasting, no gravity\n- **Open CDM**: $(\\Omega_M, \\Omega_\\Lambda) = (0.3, 0)$ -- matter only, curved\n\n\n\n```ocaml\nlet h0 = 70.0\n\nlet p_lcdm = Cosmo.lcdm ~h0 ~omega_m:0.3 ~omega_l:0.7\nlet p_edsit = Cosmo.lcdm ~h0 ~omega_m:1.0 ~omega_l:0.0\nlet p_empty = Cosmo.lcdm ~h0 ~omega_m:0.01 ~omega_l:0.0\nlet p_open = Cosmo.lcdm ~h0 ~omega_m:0.3 ~omega_l:0.0\n\nlet n_grid = 200\nlet z_grid = Nx.logspace f64 (-2.5) 0.4 n_grid (* z ~ 0.003 to 2.5 *)\n\nlet mu_of_model p =\n Nx.init f64 [| n_grid |] (fun idx ->\n let z = Nx.scalar f64 (Nx.item [idx.(0)] z_grid) in\n Nx.item [] (Cosmo.distance_modulus ~p z))\n\nlet mu_lcdm = mu_of_model p_lcdm\nlet mu_edsit = mu_of_model p_edsit\nlet mu_empty = mu_of_model p_empty\nlet mu_open = mu_of_model p_open\n\nlet () = Printf.printf \"Theory curves computed for %d redshift points\\n\" n_grid\n```\n\n\nTheory curves computed for 200 redshift points\nval h0 : float = 70.\nval p_lcdm : Umbra.Cosmo.params = \nval p_edsit : Umbra.Cosmo.params = \nval p_empty : Umbra.Cosmo.params = \nval p_open : Umbra.Cosmo.params = \nval n_grid : int = 200\nval z_grid : (float, Nx.float64_elt) Nx.t = float64 [200] \n [0.00316228, 0.00327019, ..., 2.429, 2.51189]\nval mu_of_model : Umbra.Cosmo.params -> (float, Nx.float64_elt) Nx.t = \nval mu_lcdm : (float, Nx.float64_elt) Nx.t = float64 [200] \n [30.6639, 30.737, ..., 46.4718, 46.5603]\nval mu_edsit : (float, Nx.float64_elt) Nx.t = float64 [200] \n [30.6603, 30.7332, ..., 45.6533, 45.7352]\nval mu_empty : (float, Nx.float64_elt) Nx.t = float64 [200] \n [30.662, 30.735, ..., 46.7919, 46.9042]\nval mu_open : (float, Nx.float64_elt) Nx.t = float64 [200] \n [30.6615, 30.7345, ..., 46.3264, 46.4235]\n\n\n## The Hubble diagram\n\nThe Hubble diagram plots distance modulus $\\mu$ against redshift $z$. Distant\nsupernovae ($z > 0.3$) are systematically fainter than predicted by\ndecelerating models (Einstein--de Sitter, Open CDM), showing that the expansion\nhas been **accelerating**.\n\n\n\n```ocaml\nlet to32 t = Nx.astype f32 t\n\nlet _fig =\n Hugin.layers [\n Hugin.point ~x:(to32 sn_z) ~y:(to32 sn_mu)\n ~color:(Hugin.Color.with_alpha 0.3 Hugin.Color.blue)\n ~size:2.0 ~marker:Hugin.Circle () ;\n Hugin.line ~x:(to32 z_grid) ~y:(to32 mu_lcdm)\n ~color:Hugin.Color.vermillion ~line_width:2.5\n ~label:\"ΛCDM (0.3, 0.7)\" () ;\n Hugin.line ~x:(to32 z_grid) ~y:(to32 mu_edsit)\n ~color:Hugin.Color.sky_blue ~line_width:2.0\n ~line_style:`Dashed ~label:\"EdS (1, 0)\" () ;\n Hugin.line ~x:(to32 z_grid) ~y:(to32 mu_empty)\n ~color:Hugin.Color.green ~line_width:2.0\n ~line_style:`Dotted ~label:\"Empty (0, 0)\" () ;\n Hugin.line ~x:(to32 z_grid) ~y:(to32 mu_open)\n ~color:Hugin.Color.orange ~line_width:2.0\n ~line_style:`Dash_dot ~label:\"Open (0.3, 0)\" () ;\n ]\n |> Hugin.xscale `Log\n |> Hugin.xlim 0.01 2.5\n |> Hugin.xlabel \"Redshift z\"\n |> Hugin.ylabel \"Distance modulus μ (mag)\"\n |> Hugin.title \"SN Ia Hubble Diagram (Pantheon+, 1701 light curves)\"\n |> Hugin.legend ~loc:Hugin.Lower_right\n |> Hugin.grid_lines true\n```\n\n\nval to32 : ('a, 'b) Nx.t -> (float, Nx.float32_elt) Nx.t = \nval _fig : Hugin.t =\n \n\n\nz res bin \n0.01016 -0.424629577679 0. 0.01017 -0.287976550182 0. 0.01017 -0.238776550182 0. 0.01026 0.112394684341 0. 0.01026 -0.000105315659297 0. 0.01028 0.0491444208064 0. 0.01042 -0.135579053456 0. 0.01044 -0.17646444435 0. 0.01061 0.0819784530003 0. 0.01061 -0.0055215469997 0. 0.01073 -0.215072175983 0. 0.01079 -0.267045255969 0. 0.01079 -0.154445255969 0. 0.01096 0.201426552503 0. 0.01114 0.328659998161 0. 0.01114 0.231859998161 0. 0.01122 0.414035730767 0. 0.01122 0.0129357307671 0. 0.01122 0.0740357307671 0. 0.01155 0.0993356408497 0. … … … \n
\n1590 rows × 3 columns
\n\n\nval groups : (int * Talon.t) list =\n [(0,\n \n\n\nz res bin \n0.01016 -0.424629577679 0. 0.01017 -0.287976550182 0. 0.01017 -0.238776550182 0. 0.01026 0.112394684341 0. 0.01026 -0.000105315659297 0. 0.01028 0.0491444208064 0. 0.01042 -0.135579053456 0. 0.01044 -0.17646444435 0. 0.01061 0.0819784530003 0. 0.01061 -0.0055215469997 0. 0.01073 -0.215072175983 0. 0.01079 -0.267045255969 0. 0.01079 -0.154445255969 0. 0.01096 0.201426552503 0. 0.01114 0.328659998161 0. 0.01114 0.231859998161 0. 0.01122 0.414035730767 0. 0.01122 0.0129357307671 0. 0.01122 0.0740357307671 0. 0.01155 0.0993356408497 0. … … … \n
\n24 rows × 3 columns
\n\n);\n (1,\n \n\n\nz res bin \n0.01246 0.178978224041 1. 0.01258 0.0374364117569 1. 0.01258 0.0540364117569 1. 0.01259 -0.084899767747 1. 0.01259 -0.154499767747 1. 0.01279 0.0424614815497 1. 0.01283 -0.0770620111033 1. 0.01283 -0.210762011103 1. 0.01303 -0.0118654606603 1. 0.01303 0.0133345393397 1. 0.01304 -0.173842071152 1. 0.01304 -0.181742071152 1. 0.01312 -0.242409144024 1. 0.01312 -0.290709144024 1. 0.01325 0.184241133384 1. 0.01325 0.200641133384 1. 0.01325 0.140741133384 1. 0.01375 -0.0679294440597 1. 0.01375 -0.0966294440597 1. 0.01375 -0.0929294440597 1. … … … \n
\n52 rows × 3 columns
\n\n);\n (2,\n \n\n\nz res bin \n0.01546 -0.0881971344432 2. 0.01549 0.14766106801 2. 0.0155 0.131148947175 2. 0.0155 0.0842489471753 2. 0.0155 -0.127551052825 2. 0.0155 -0.241751052825 2. 0.01557 -0.346610654758 2. 0.01557 -0.329610654758 2. 0.01562 0.0293736691777 2. 0.01562 -0.0136263308223 2. 0.01565 0.183674953401 2. 0.01576 0.00234770294004 2. 0.01578 -0.16752766025 2. 0.01578 -0.0964276602495 2. 0.01581 -0.177684167024 2. 0.01581 -0.269884167024 2. 0.01587 0.213026247358 2. 0.01588 -0.151252326051 2. 0.0159 -0.0164068904934 2. 0.0159 -0.0504068904934 2. … … … \n
\n72 rows × 3 columns
\n\n);\n (3,\n \n\n\nz res bin \n0.01947 0.148326584074 3. 0.01947 0.182126584074 3. 0.01975 0.0993213367396 3. 0.01975 0.0817213367396 3. 0.01976 0.0277114394619 3. 0.01995 -0.31497156999 3. 0.02001 -0.000956680793522 3. 0.02006 -0.17972935267 3. 0.02019 -0.198795323932 3. 0.02019 -0.00399532393241 3. 0.02023 -0.110835916616 3. 0.02023 -0.209135916616 3. 0.02023 -0.134335916616 3. 0.02023 -0.0621359166155 3. 0.02024 0.287380263147 3. 0.02034 0.0942711314822 3. 0.02034 -0.147328868518 3. 0.02035 -0.152406886052 3. 0.02035 -0.157706886052 3. 0.02035 -0.230206886052 3. … … … \n
\n92 rows × 3 columns
\n\n);\n (4,\n \n\n\nz res bin \n0.02388 -0.0795275874192 4. 0.0239 -0.187266827175 4. 0.0239 -0.0822668271749 4. 0.02391 -0.177885876584 4. 0.02401 0.0154444709488 4. 0.02411 0.00941249193803 4. 0.02411 -0.624387508062 4. 0.02412 -0.218198645962 4. 0.02417 -0.158848742089 4. 0.02417 -0.0169487420891 4. 0.02428 0.0693737074972 4. 0.02429 -0.0347311260491 4. 0.02432 -0.127443419316 4. 0.02432 -0.161543419316 4. 0.02432 0.117156580684 4. 0.02434 -0.106149778375 4. 0.02453 -0.125437393885 4. 0.02453 -0.0588373938845 4. 0.02453 -0.184137393885 4. 0.02457 0.0413818842716 4. … … … \n
\n112 rows × 3 columns
\n\n);\n (5,\n \n\n\nz res bin \n0.02969 -0.275399367382 5. 0.02978 -0.058667628548 5. 0.02978 -0.048567628548 5. 0.0299 -0.133627805804 5. 0.02996 0.287555242197 5. 0.03012 -0.0656808035682 5. 0.03012 -0.101380803568 5. 0.03012 -0.109080803568 5. 0.03023 -0.0747137430286 5. 0.03031 -0.0448378058181 5. 0.03036 -0.122370156394 5. 0.03047 -0.0761406185614 5. 0.03059 -0.260703391037 5. 0.03075 -0.278601806314 5. 0.03076 -0.0231184984297 5. 0.03083 -0.171728924107 5. 0.03086 -0.222272818751 5. 0.03091 -0.149641417105 5. 0.03096 0.00229566781329 5. 0.03108 0.26806774982 5. … … … \n
\n99 rows × 3 columns
\n\n);\n (6,\n \n\n\nz res bin \n0.03697 0.037370558993 6. 0.03702 -0.234817279977 6. 0.03702 -0.111717279977 6. 0.03702 0.0290827200227 6. 0.03702 0.0395827200227 6. 0.03705 -0.0785080806863 6. 0.03707 0.00409884352897 6. 0.03725 -0.047410467157 6. 0.03725 -0.061410467157 6. 0.0373 -0.187276253133 6. 0.0374 -0.0840961258395 6. 0.03753 -0.121668755977 6. 0.03756 0.260864344984 6. 0.03756 0.146064344984 6. 0.03787 -0.0843128672071 6. 0.0379 -0.0299641891915 6. 0.03796 -0.00696275219823 6. 0.03818 -0.214844515711 6. 0.03818 -0.246944515711 6. 0.03828 -0.00033051522886 6. … … … \n
\n61 rows × 3 columns
\n\n);\n (7,\n \n\n\nz res bin \n0.0459 -0.135594053043 7. 0.04625 -0.269958777154 7. 0.04631 0.00296267259202 7. 0.04643 -0.217483496434 7. 0.04656 -0.0402921381081 7. 0.04664 -0.0809044161309 7. 0.04682 -0.0821587037548 7. 0.04682 -0.0750587037548 7. 0.04691 0.193176209883 7. 0.04738 -0.0229677846992 7. 0.0476 0.00344066051645 7. 0.04777 -0.010780094188 7. 0.04777 -0.042680094188 7. 0.04819 -0.127231460348 7. 0.04837 -0.259617069447 7. 0.0486 -0.313360481081 7. 0.04865 -0.00494607201814 7. 0.04934 -0.0775548977512 7. 0.0494 -0.21085715033 7. 0.04944 -0.0989568708902 7. … … … \n
\n47 rows × 3 columns
\n\n);\n (8,\n \n\n\nz res bin \n0.05708 -0.0115206686667 8. 0.05728 0.0590741402491 8. 0.05824 -0.064625197056 8. 0.05824 -0.114325197056 8. 0.0583 -0.0945240955029 8. 0.05886 -0.0901701116884 8. 0.05886 0.0756298883116 8. 0.05974 -0.803417802658 8. 0.06092 -0.115228066137 8. 0.06099 0.0133048886582 8. 0.06099 0.0427048886582 8. 0.06121 -0.0195443596127 8. 0.06137 -0.202880711986 8. 0.06137 -0.194580711986 8. 0.06153 -0.0948022912486 8. 0.06372 0.103960473042 8. 0.06384 -0.232250654126 8. 0.06446 -0.194986434258 8. 0.06533 -0.218508110773 8. 0.06627 0.0260876633256 8. … … … \n
\n32 rows × 3 columns
\n\n);\n (9,\n \n\n\nz res bin \n0.07089 -0.0324754668871 9. 0.0709 0.242007811196 9. 0.07091 -0.116208867472 9. 0.07116 -0.140111813817 9. 0.07158 -0.259128451072 9. 0.07167 -0.0958508194342 9. 0.07193 0.0423148998275 9. 0.07222 0.103375550607 9. 0.07252 0.0846613871215 9. 0.07393 -0.0276218037538 9. 0.0744 -0.16077227092 9. 0.07446 -0.147085210939 9. 0.0752 -0.143129423208 9. 0.0752 -0.0538294232081 9. 0.0756 0.182834610902 9. 0.07575 -0.0897256482245 9. 0.07588 -0.00548430816918 9. 0.07845 -0.143184159753 9. 0.07859 0.112798691001 9. 0.07875 0.292116111885 9. … … … \n
\n33 rows × 3 columns
\n\n);\n (10,\n \n\n\nz res bin \n0.0887 -0.214759965462 10. 0.09039 -0.114490123038 10. 0.09089 -0.0202850975671 10. 0.09205 0.0842789146355 10. 0.09293 0.125110163521 10. 0.0995 0.185107403999 10. 0.10165 -0.00302391994575 10. 0.10221 0.0947708510874 10. 0.10246 -0.0681907017381 10. 0.10294 -0.148432610935 10. 0.10361 -0.0109079016729 10. 0.10374 -0.139664168122 10. 0.10507 0.0529089143211 10. 0.10661 -0.148665966028 10. 0.10707 0.0428133675802 10. 0.10711 0.00906130066566 10. 0.10713 0.359235381084 10. 0.10774 0.064581151469 10. 0.10794 -0.0363509057905 10. 0.10908 -0.0604317214155 10. \n
\n20 rows × 3 columns
\n\n);\n (11,\n \n\n\nz res bin \n0.11001 0.0886813554461 11. 0.11259 0.0870049833038 11. 0.11388 -0.172851036977 11. 0.1165 -0.0717173699591 11. 0.11653 0.0170929257421 11. 0.1176 -0.0613460214423 11. 0.11792 -0.031472958122 11. 0.11818 -0.229620528577 11. 0.11901 -0.0984636003687 11. 0.12014 -0.0872353293955 11. 0.12058 0.101178454287 11. 0.12086 -0.0619431118695 11. 0.12207 -0.125106114742 11. 0.12231 -0.112215348847 11. 0.12278 -0.0476216623284 11. 0.12316 -0.143318307083 11. 0.12357 -0.13125195377 11. 0.12377 0.0495330304242 11. 0.12383 0.0717196359588 11. 0.12393 -0.103934884938 11. … … … \n
\n38 rows × 3 columns
\n\n);\n (12,\n \n\n\nz res bin \n0.13614 0.00634512405369 12. 0.13658 -0.130706237119 12. 0.137 0.0242022119563 12. 0.13713 -0.0258886328085 12. 0.13745 0.35822685965 12. 0.13822 -0.12008128119 12. 0.13826 -0.0561499783706 12. 0.1384 0.0141110181508 12. 0.13851 0.0294747952992 12. 0.13875 -0.0147267381425 12. 0.1388 0.0312404310394 12. 0.13955 -0.0066181867595 12. 0.14082 0.128628511143 12. 0.14104 0.231016921037 12. 0.14123 -0.0148979084376 12. 0.14134 0.0663005741243 12. 0.14325 -0.0414714699277 12. 0.14345 0.423697521 12. 0.14359 0.0195383381869 12. 0.14404 -0.0256092965459 12. … … … \n
\n64 rows × 3 columns
\n\n);\n (13,\n \n\n\nz res bin \n0.16924 0.0389733472685 13. 0.16971 0.189083761363 13. 0.17042 -0.0185879074946 13. 0.17124 0.0383736768909 13. 0.17169 -0.0498724221099 13. 0.17256 -0.0199123832614 13. 0.1727 -0.25631246117 13. 0.17297 0.19872715637 13. 0.17331 -0.145074641558 13. 0.17374 0.00231747862402 13. 0.17378 -0.189222107681 13. 0.17392 0.0581902517402 13. 0.17417 0.0296229858146 13. 0.1742 -0.212480783249 13. 0.17438 -0.0413020372001 13. 0.17443 0.00732580573651 13. 0.17444 -0.062508604123 13. 0.17498 0.188743907993 13. 0.17666 -0.0606712840265 13. 0.17713 -0.0738066505178 13. … … … \n
\n117 rows × 3 columns
\n\n);\n (14,\n \n\n\nz res bin \n0.21037 -0.53515754654 14. 0.21084 -0.031562219736 14. 0.21095 -0.190002164661 14. 0.21114 -0.0429424845705 14. 0.21134 0.14110646419 14. 0.21174 -0.0372897541268 14. 0.212 -0.0942081001471 14. 0.21225 -0.0732110933607 14. 0.2135 -0.00278059242427 14. 0.21365 -0.0178518660398 14. 0.21398 -0.142524867022 14. 0.2144 -0.0166920577869 14. 0.21507 -0.159519939026 14. 0.21521 0.420530657906 14. 0.21578 0.175431938785 14. 0.2165 -0.232502482431 14. 0.21689 0.00510984041876 14. 0.21692 0.0307803130448 14. 0.21742 0.0715943550243 14. 0.21794 -0.0793987415813 14. … … … \n
\n142 rows × 3 columns
\n\n);\n (15,\n \n\n\nz res bin \n0.26141 -0.063389305081 15. 0.26162 -0.0941332792955 15. 0.26172 0.174141517223 15. 0.26173 0.16764901455 15. 0.26175 -0.121535981157 15. 0.26184 0.179131697137 15. 0.262 0.205752656004 15. 0.26303 0.738750935002 15. 0.26323 0.111609861946 15. 0.2636 -0.130892774852 15. 0.26393 0.0251761002073 15. 0.26397 -0.230491074865 15. 0.26408 0.0991994542632 15. 0.26419 -0.0286096346745 15. 0.2646 -0.0244674248267 15. 0.26463 0.131857822634 15. 0.26582 0.0577820584371 15. 0.26583 -0.226609147044 15. 0.2664 -0.266102716605 15. 0.267 -0.0153587433182 15. … … … \n
\n150 rows × 3 columns
\n\n);\n (16,\n \n\n\nz res bin \n0.32548 0.180663773592 16. 0.3256 0.214952098223 16. 0.3258 0.114833307567 16. 0.32581 -0.00134261005185 16. 0.32632 -0.00261164522524 16. 0.32804 -0.0658203680861 16. 0.32842 0.23901384648 16. 0.32848 0.281661625296 16. 0.32851 -0.135564457581 16. 0.32868 0.000754754949696 16. 0.32868 0.0291547549497 16. 0.32871 -0.143471204838 16. 0.32907 -0.017181284091 16. 0.32941 0.334561631049 16. 0.32952 0.161934844394 16. 0.32968 0.0247326862566 16. 0.32995 -0.235994772669 16. 0.33047 0.146304669334 16. 0.33056 0.082430130099 16. 0.33063 0.0834056020207 16. … … … \n
\n130 rows × 3 columns
\n\n);\n (17,\n \n\n\nz res bin \n0.40368 0.0305227159842 17. 0.40463 0.0832671563379 17. 0.40483 -0.255385075024 17. 0.4055 0.183823918891 17. 0.40646 0.0885295187297 17. 0.40895 0.0818394810777 17. 0.4092 0.120588849089 17. 0.40935 0.0852588703231 17. 0.40949 0.0566911608651 17. 0.41004 0.0584848294338 17. 0.41123 0.0295285142735 17. 0.4114 0.199979142478 17. 0.41161 0.178283386575 17. 0.41266 0.148013322853 17. 0.41657 0.0642467660022 17. 0.41857 -0.0600359425703 17. 0.41936 0.00406598556371 17. 0.41939 -0.222016063038 17. 0.4196 0.182509917202 17. 0.41965 0.114106661782 17. … … … \n
\n97 rows × 3 columns
\n\n);\n (18,\n \n\n\nz res bin \n0.50282 0.0115974630251 18. 0.50285 -0.0792578980026 18. 0.50306 -0.00224520003424 18. 0.50316 -0.00676282447449 18. 0.50593 -0.0864656605981 18. 0.50615 -0.0622987115987 18. 0.50725 0.0635424328218 18. 0.50739 0.0689229785269 18. 0.50825 -0.15809275328 18. 0.51016 -0.0926766588196 18. 0.51095 -0.132914123298 18. 0.51169 0.193508853341 18. 0.51387 0.0506093976964 18. 0.51437 0.0172694047828 18. 0.51469 0.11564493179 18. 0.51726 -0.130069979609 18. 0.51883 -0.0362931904153 18. 0.51885 -0.0488939890317 18. 0.51941 0.00968501476334 18. 0.51968 0.0550258324352 18. … … … \n
\n96 rows × 3 columns
\n\n);\n (19,\n \n\n\nz res bin \n0.62525 -0.0676927812368 19. 0.62725 -0.148965950221 19. 0.63077 -0.0857981161031 19. 0.63183 -0.442410797453 19. 0.63225 0.094002948885 19. 0.63399 -0.165486457732 19. 0.63777 -0.248879776508 19. 0.63794 0.0843028521275 19. 0.63824 -0.121262699059 19. 0.63873 -0.00232867361532 19. 0.63934 0.00470128976596 19. 0.64185 0.0517482101598 19. 0.64311 -0.440236071775 19. 0.64371 -0.00444929008427 19. 0.64371 -0.000249290084263 19. 0.6477 -0.00131150843364 19. 0.64852 0.215975033203 19. 0.6487 -0.0909737694835 19. 0.64962 -0.0477982164029 19. 0.66213 0.0994509027894 19. … … … \n
\n76 rows × 3 columns
\n\n);\n (20,\n \n\n\nz res bin \n0.77929 -0.286728045808 20. 0.78807 -0.0531352325948 20. 0.78907 0.0228403967112 20. 0.78928 -0.162099242024 20. 0.79662 -0.484747596675 20. 0.79863 0.039436351798 20. 0.83981 -0.174927172808 20. 0.83981 -0.333727172808 20. 0.85482 0.0320794309533 20. 0.93585 -0.308388644025 20. \n
\n10 rows × 3 columns
\n\n);\n (21,\n \n\n\nz res bin \n0.97423 0.154450661289 21. 1.01242 0.0501694564944 21. 1.01988 0.186919560833 21. 1.02088 -0.121019067405 21. 1.02789 0.42494712992 21. 1.04817 0.254697417559 21. 1.12092 0.0759845634931 21. \n
\n7 rows × 3 columns
\n\n);\n (22,\n \n\n\nz res bin \n1.23225 -0.324587237938 22. 1.23597 0.250701908425 22. 1.29911 0.0602023436219 22. 1.3041 -0.0948606170644 22. 1.30611 0.199592161811 22. 1.31317 -0.0496838731239 22. 1.3291 0.00165723209277 22. 1.34101 -0.145764204634 22. 1.35136 -0.373984554953 22. 1.35608 -0.124370205399 22. 1.39103 -0.216813113651 22. 1.41633 -0.608368865186 22. \n
\n12 rows × 3 columns
\n\n);\n (23,\n \n\n\nz res bin \n1.5429 -0.239796061037 23. 1.54901 -0.0492647551793 23. 1.61505 -0.312860625322 23. 1.69706 -0.341579608953 23. 1.80119 -0.33004923418 23. \n
\n5 rows × 3 columns
\n\n)]\nval n_groups : int = 24\nval bz : (float, Nx.float32_elt) Nx.t = float32 [24] \n [0.0109479, 0.0140023, ..., 1.32297, 1.64104]\nval bmu : (float, Nx.float32_elt) Nx.t = float32 [24] \n [0.00578864, -0.0499531, ..., -0.118857, -0.25471]\nval berr : (float, Nx.float32_elt) Nx.t = float32 [24] \n [0.043205, 0.0225362, ..., 0.070028, 0.0543295]\n\n\n\n```ocaml\nlet _fig =\n Hugin.layers [\n Hugin.errorbar ~x:bz ~y:bmu ~yerr:(`Symmetric berr)\n ~color:Hugin.Color.black ~cap_size:4.0 ~line_width:1.5 () ;\n Hugin.point ~x:bz ~y:bmu\n ~color:Hugin.Color.black ~size:5.0 ~marker:Hugin.Circle () ;\n Hugin.line ~x:(to32 z_grid) ~y:(to32 res_lcdm)\n ~color:Hugin.Color.vermillion ~line_width:2.5\n ~label:\"ΛCDM (0.3, 0.7)\" () ;\n Hugin.line ~x:(to32 z_grid) ~y:(to32 res_edsit)\n ~color:Hugin.Color.sky_blue ~line_width:2.0\n ~line_style:`Dashed ~label:\"EdS (1, 0)\" () ;\n Hugin.line ~x:(to32 z_grid) ~y:(to32 res_open)\n ~color:Hugin.Color.orange ~line_width:2.0\n ~line_style:`Dash_dot ~label:\"Open (0.3, 0)\" () ;\n Hugin.hline ~y:0.0 ~line_style:`Dotted ~color:Hugin.Color.gray () ;\n ]\n |> Hugin.xscale `Log\n |> Hugin.xlim 0.01 2.5\n |> Hugin.xlabel \"Redshift z\"\n |> Hugin.ylabel \"Δμ (mag, relative to empty universe)\"\n |> Hugin.title \"Hubble Residuals: The Acceleration Signal\"\n |> Hugin.legend ~loc:Hugin.Upper_left\n |> Hugin.grid_lines true\n```\n\n\nval _fig : Hugin.t =\n \n\n\nCID IDSURVEY zHD zHDERR zCMB zCMBERR zHEL zHELERR m_b_corr m_b_corr_err_DIAG … \n2013E 56 0.01016 0.00085 0.01042 8e-05 0.00936 8e-05 13.5264 0.3475 … 1999ac 57 0.01017 0.00084 0.00979 2e-05 0.00947 2e-05 13.6652 0.364224 … 1999ac 62 0.01017 0.00084 0.00979 2e-05 0.00947 2e-05 13.7144 0.34081 … 2009an 51 0.01026 0.00084 0.00921 1e-05 0.00887 1e-05 14.0848 0.305101 … 2009an 65 0.01026 0.00084 0.00921 1e-05 0.00887 1e-05 13.9723 0.297865 … 2006bh 5 0.01028 0.00086 0.01042 0.00015 0.01077 0.00015 14.0258 0.246478 … 2004S 57 0.01042 0.00084 0.0098 2e-05 0.0093 2e-05 13.8706 0.316076 … 2021hpr 57 0.01044 0.00084 0.00958 2e-05 0.00938 2e-05 13.8339 0.342855 … 2002dp 63 0.01061 0.00084 0.01049 1e-05 0.01169 1e-05 14.1276 0.307827 … 2002dp 57 0.01061 0.00084 0.01049 1e-05 0.01169 1e-05 14.0401 0.273239 … 1997do 62 0.01073 0.00084 0.01048 2e-05 0.01012 2e-05 13.8551 0.363667 … 1997bq 62 0.01079 0.00084 0.00993 2e-05 0.00973 2e-05 13.8153 0.322889 … 2008fv_comb 50 0.01079 0.00084 0.00993 2e-05 0.00973 2e-05 13.9279 0.377003 … ASASSN-16jf 150 0.01096 0.00084 0.0104 1e-05 0.01144 1e-05 14.3179 0.313698 … iPTF13ebh 56 0.01114 0.00085 0.01238 5e-05 0.01317 5e-05 14.4807 0.341421 … iPTF13ebh 5 0.01114 0.00085 0.01238 5e-05 0.01317 5e-05 14.3839 0.293983 … 2010ko 56 0.01122 0.00084 0.01096 2e-05 0.01082 2e-05 14.5817 0.352878 … 2013ex 51 0.01122 0.00084 0.01096 2e-05 0.01082 2e-05 14.1806 0.282135 … 2013ex 56 0.01122 0.00084 0.01096 2e-05 0.01082 2e-05 14.2417 0.32405 … 2009ab 5 0.01155 0.00085 0.01189 8e-05 0.01219 8e-05 14.3303 0.27987 … … … … … … … … … … … … \n
\n1590 rows × 47 columns
\n\n\nval hf_col : string -> (float, Nx.float64_elt) Nx.t = \nval hf_z : (float, Nx.float64_elt) Nx.t = float64 [1590] \n [0.01016, 0.01017, ..., 1.91165, 2.26137]\nval hf_mu : (float, Nx.float64_elt) Nx.t = float64 [1590] \n [32.7794, 32.9182, ..., 45.4233, 46.1828]\nval hf_w : (float, Nx.float64_elt) Nx.t = float64 [1590] \n [8.23147, 7.49694, ..., 7.77459, 12.6367]\nval n_hf : int = 1590\nval hf_z_arr : float array =\n [|0.01016; 0.01017; 0.01017; 0.01026; 0.01026; 0.01028; 0.01042; 0.01044;\n 0.01061; 0.01061; 0.01073; 0.01079; 0.01079; 0.01096; 0.01114; 0.01114;\n 0.01122; 0.01122; 0.01122; 0.01155; 0.01195; 0.01213; 0.0122; 0.01233;\n 0.01246; 0.01258; 0.01258; 0.01259; 0.01259; 0.01279; 0.01283; 0.01283;\n 0.01303; 0.01303; 0.01304; 0.01304; 0.01312; 0.01312; 0.01325; 0.01325;\n 0.01325; 0.01375; 0.01375; 0.01375; 0.01376; 0.01386; 0.01388; 0.01389;\n 0.01389; 0.01411; 0.01424; 0.01442; 0.01442; 0.01442; 0.01442; 0.01442;\n 0.01446; 0.0145; 0.01453; 0.0146; 0.01462; 0.01463; 0.01463; 0.01467;\n 0.01472; 0.01484; 0.01492; 0.01493; 0.01499; 0.01499; 0.01515; 0.01519;\n 0.01525; 0.01529; 0.01542; 0.01543; 0.01546; 0.01549; 0.0155; 0.0155;\n 0.0155; 0.0155; 0.01557; 0.01557; 0.01562; 0.01562; 0.01565; 0.01576;\n 0.01578; 0.01578; 0.01581; 0.01581; 0.01587; 0.01588; 0.0159; 0.0159;\n 0.0159; 0.01603; 0.01652; 0.01652; 0.01656; 0.01657; 0.01662; 0.01666;\n 0.01671; 0.01678; 0.01682; 0.01682; 0.01682; 0.0169; 0.0169; 0.01692;\n 0.01698; 0.01699; 0.01705; 0.01718; 0.01718; 0.0172; 0.0173; 0.0173;\n 0.01733; 0.01733; 0.01734; 0.01737; 0.01737; 0.01743; 0.01747; 0.01747;\n 0.01752; 0.01776; 0.01778; 0.01778; 0.01784; 0.01784; 0.0179; 0.01802;\n 0.01808; 0.01826; 0.01826; 0.01839; 0.01855; 0.01855; 0.01865; 0.01865;\n 0.01866; 0.01875; 0.01875; 0.01905; 0.01947; 0.01947; 0.01975; 0.01975;\n 0.01976; 0.01995; 0.02001; 0.02006; 0.02019; 0.02019; 0.02023; 0.02023;\n 0.02023; 0.02023; 0.02024; 0.02034; 0.02034; 0.02035; 0.02035; 0.02035;\n 0.02044; 0.02049; 0.02052; 0.02056; 0.02056; 0.02081; 0.02082; 0.02082;\n 0.0209; 0.02096; 0.02106; 0.02116; 0.02116; 0.02118; 0.02118; 0.02131;\n 0.02131; 0.02131; 0.02134; 0.02137; 0.02151; 0.02153; 0.0217; 0.02183;\n 0.02183; 0.02197; 0.02198; 0.02203; 0.02205; 0.02207; 0.02215; 0.02219;\n 0.02228; 0.02228; 0.02231; 0.02234; 0.02234; 0.02236; 0.02239; 0.02239;\n 0.0224; 0.0224; 0.02241; 0.02255; 0.02266; 0.0227; 0.02273; 0.02295;\n 0.02295; 0.02298; 0.02298; 0.02303; 0.02307; 0.02313; 0.02316; 0.02321;\n 0.02325; 0.02331; 0.02331; 0.02342; 0.02342; 0.02342; 0.02343; 0.02343;\n 0.02344; 0.02352; 0.02354; 0.02357; 0.02357; 0.02357; 0.02365; 0.02369;\n 0.02388; 0.0239; 0.0239; 0.02391; 0.02401; 0.02411; 0.02411; 0.02412;\n 0.02417; 0.02417; 0.02428; 0.02429; 0.02432; 0.02432; 0.02432; 0.02434;\n 0.02453; 0.02453; 0.02453; 0.02457; 0.02462; 0.02462; 0.02462; 0.02464;\n 0.02464; 0.02466; 0.02491; 0.02494; 0.02509; 0.0251; 0.0251; 0.0251;\n 0.0251; 0.02512; 0.02513; 0.02517; 0.02517; 0.02517; 0.02519; 0.02519;\n 0.02521; 0.02525; 0.02525; 0.02534; 0.02534; 0.02556; 0.02557; 0.02585;\n 0.02591; 0.02596; 0.02598; 0.02598; 0.02598; 0.02598; 0.02626; 0.02626;\n 0.02632; 0.02669; 0.02691; ...|]\nval hf_mu_arr : float array =\n [|32.7794; 32.9182; 32.9674; 33.3378; 33.2253; 33.2788; 33.1236; 33.0869;\n 33.3806; 33.2931; 33.1081; 33.0683; 33.1809; 33.5709; 33.7337; 33.6369;\n 33.8347; 33.4336; 33.4947; 33.5833; 33.4974; 33.8783; 33.7023; 33.7403;\n 33.8286; 33.708; 33.7246; 33.5874; 33.5178; 33.7492; 33.6365; 33.5028;\n 33.7355; 33.7607; 33.5752; 33.5673; 33.52; 33.4717; 33.9682; 33.9846;\n 33.9247; 33.797; 33.7683; 33.772; 33.7649; 33.6741; 33.6968; 33.8053;\n 33.6651; 33.8206; 33.8444; 34.0187; 34.0249; 34.3366; 34.1781; 34.252;\n 33.9078; 34.1727; 33.7526; 33.8826; 33.9253; 33.7553; 33.7189; 33.8916;\n 33.7894; 34.0109; 34.1902; 33.8524; 33.974; 33.9261; 33.8877; 34.1793;\n 33.7706; 34.0334; 34.3426; 34.0744; 34.0331; 34.2732; 34.2581; 34.2112;\n 33.9994; 33.8852; 33.7902; 33.8072; 34.1732; 34.1302; 34.3317; 34.1657;\n 33.9986; 34.0697; 33.9926; 33.9004; 34.3916; 34.0287; 34.1663; 34.1323;\n 34.0329; 34.1062; 34.4137; 33.9999; 34.1608; 33.9726; 34.2793; 34.3408;\n 33.8985; 34.1679; 34.4649; 34.5341; 33.9673; 33.841; 33.8975; 34.4894;\n 34.3409; 34.0755; 34.1973; 34.022; 34.2003; 34.6084; 34.1436; 34.1517;\n 34.4009; 34.5101; 34.361; 34.31; 34.2669; 33.7799; 34.3228; 34.3133;\n 34.1465; 34.2221; 34.3364; 34.2464; 34.2939; 34.2418; 34.2975; 34.2661;\n 34.6078; 34.9792; 34.2005; 34.4393; 34.492; 34.576; 34.6171; 34.5716;\n 34.6066; 34.5466; 34.4866; 34.4286; 34.7747; 34.8085; 34.757; 34.7394;\n 34.6865; 34.3648; 34.6854; 34.5121; 34.5072; 34.702; 34.5995; 34.5012;\n 34.576; 34.6482; 34.9988; 34.8165; 34.5749; 34.5709; 34.5656; 34.4931;\n 34.6552; 35.1263; 34.556; 34.4168; 34.5647; 34.4863; 34.674; 34.6349;\n 34.8118; 34.6956; 34.6693; 34.375; 34.8194; 35.118; 35.2056; 34.7561;\n 34.7344; 34.7419; 34.7627; 34.3468; 34.5899; 34.6114; 34.7709; 34.7583;\n 34.7742; 34.8589; 34.9654; 34.6739; 34.7477; 34.8847; 35.0481; 35.0136;\n 34.8059; 34.8234; 34.9245; 34.7727; 34.8704; 34.7814; 34.7714; 34.8603;\n 34.759; 34.7997; 35.0392; 34.7306; 34.8438; 34.8693; 34.6849; 35.2685;\n 34.8584; 34.9992; 34.9712; 34.985; 34.7666; 35.0225; 34.883; 34.9168;\n 35.0537; 34.6478; 34.4301; 34.9851; 35.0081; 34.9783; 34.8582; 34.8582;\n 34.8331; 34.8391; 34.9998; 35.0406; 34.9509; 34.5715; 34.9543; 34.943;\n 34.9949; 34.889; 34.994; 34.8993; 35.1018; 35.1049; 34.4711; 34.8782;\n 34.9421; 35.084; 35.1803; 35.0771; 34.9871; 34.953; 35.2317; 35.0102;\n 35.008; 35.0746; 34.9493; 35.1784; 35.1061; 35.0521; 34.8898; 34.9528;\n 34.807; 34.8112; 35.0862; 35.1973; 35.2152; 34.8369; 34.8482; 35.079;\n 35.009; 35.0276; 34.4583; 34.9117; 35.1989; 35.2277; 35.1139; 35.1145;\n 34.9842; 34.9441; 34.9063; 35.1882; 35.1974; 35.1828; 35.0621; 34.7664;\n 35.1181; 35.4014; 35.2931; 35.1081; 34.835; 35.3111; 35.2855; 35.3429;\n 35.0223; 35.308; 35.2688; ...|]\nval hf_w_arr : float array =\n [|8.23146814613716593; 7.4969352680999215; 8.55574221326686413;\n 10.6591609759993187; 11.1791254824205311; 16.2654007737080022;\n 9.93710117748064903; 8.45464229847009463; 10.4726778114779471;\n 13.2645900552113254; 7.5197722886281948; 9.52504612371758519;\n 6.99983232092770624; 10.0871977557366836; 8.52534214273321567;\n 11.4738309842553772; 7.98393120275620927; 12.4487789115859133;\n 9.45744263313261868; 12.6492602333561486; 11.0006818657807397;\n 8.3347365445043; 14.0625703127636719; 15.4835677155751181;\n 14.227591884866; 17.2516774910131865; 15.987591222375503;\n 14.7860742788850779; 14.1129085551366398; 5.76023962347622742;\n 12.1084887207242176; 9.1205006900497132; 17.3925484174557248;\n 12.1218985292262946; 11.683103651051935; 12.3299901204144131;\n 13.365063188000093; 14.6123881673806135; 10.8425251089593679;\n 11.1557637660110593; 10.7776256002286637; 12.9527972436885257;\n 18.4550886236199396; 17.9919475158734201; 11.9903729275141533;\n 8.38375008416108436; 18.3362669046649529; 11.9893765286826248;\n 11.3974938041915532; 15.8888386016909919; 14.1875346969754474;\n 10.5967952266739616; 8.05528722258021723; 7.27645222224196;\n 6.34775702674259801; 6.33731050419107; 17.2868421478877323;\n 7.50490609740132886; 21.3740568964570166; 16.0879126286809289;\n 16.8064771777973228; 6.55759480278037898; 6.48056420082261653;\n 17.2611398525342; 6.33357902093675751; 13.69489137836125;\n 7.20015441460373751; 12.9805322684816744; 18.9754652580740206;\n 18.625928986051818; 11.4921195365949078; 14.0516078939454445;\n 10.9506480889754076; 8.43554877406980452; 7.19189246515212055;\n 10.0869414622828693; 19.4983479754138251; 8.63222933841724327;\n 13.4641128131653112; 15.0148704932793411; 12.1982247035558142;\n 10.146793599690735; 11.8434308491135152; 11.8582808284653058;\n 14.3086498722222988; 8.94749089006965; 12.873171291985976;\n 12.1506484441415417; 20.672995317749784; 11.8593426098531527;\n 17.4644346112094517; 20.851416708835508; 13.6828374657754281;\n 19.7840253229750367; 16.9500116927867452; 22.7904126118422568;\n 13.4657927237540278; 12.6121817743929938; 8.3352659411483927;\n 8.9251047074459; 21.0194316349305446; 5.79853079888242284;\n 14.5772613757875487; 25.5016762892044468; 16.6077598589186302;\n 16.4609255025734029; 7.78826881381891045; 8.56872028090763749;\n 12.362786901820888; 11.9473092996875536; 14.2871324174975616;\n 10.3847790563541444; 19.1715506082083955; 11.2747485236470109;\n 21.741787484924366; 10.4736268288461609; 9.44831673156642182;\n 20.6961375184077028; 12.5135089429374613; 15.1242634067501207;\n 20.3458786329970458; 18.7015599950648692; 20.6077296866621609;\n 14.3963000301706341; 14.520768051703806; 9.14760924209767445;\n 24.9398587271691525; 25.5540420407535684; 18.6728015422666118;\n 9.74052699653846; 22.6509220160585194; 21.9531684900774628;\n 16.4439751191577166; 20.1324567073815182; 9.83237146373930848;\n 18.2299477227476743; 24.6365123425761; 20.9496429577490275;\n 12.7035098104889101; 19.4234861464175346; 15.3351298538964578;\n 11.3947238832405269; 17.509040539551691; 17.3627026557141519;\n 12.988673539837988; 21.5302476249783794; 9.70348417833172761;\n 12.6278730549272478; 26.1901234656971589; 26.1558451309293041;\n 4.59293241178748; 5.32095686553453; 6.05944007446558164;\n 11.4346774646612896; 23.9645656337853978; 25.1289970704292536;\n 18.27955215083119; 26.9972194699756898; 28.2005673569440063;\n 11.4469833641548; 27.7291367056932039; 28.3439800971673748;\n 28.3793238806152281; 17.1000589321468475; 7.0800382049997026;\n 28.3814405615375804; 31.6626267020978887; 27.4834283515611411;\n 21.7158579836266306; 7.08150786072332394; 8.96926295523383743;\n 21.2058816479978; 16.0635486772255796; 9.15802089756929;\n 18.7566768373479498; 15.4378524337368379; 11.0946109925032541;\n 21.335767528888681; 18.7911672303113519; 24.2452796480947299;\n 25.9182978468908658; 10.3599255061832096; 10.8133084385166747;\n 30.6648936516269401; 31.8555650038301259; 28.3678374577575063;\n 8.65199489673906541; 25.2493498917324288; 25.2070269959166922;\n 6.46396662873246; 26.0672487054426796; 12.4466708843980065;\n 15.5433263008844627; 16.4646661210180483; 21.2416677860634415;\n 16.2873329914842024; 25.2277949897877924; 28.8960644281234629;\n 12.0035869468058891; 21.4019504835921; 30.8025616050924036;\n 15.0103333911976673; 34.6685262497122935; 25.2508724579438208;\n 23.9167726307908026; 19.8123910920625654; 10.9170968668811543;\n 16.9183741772662302; 11.9963539297717094; 15.3461854905195416;\n 21.7290195082093263; 16.9154518357872767; 20.5587966414740855;\n 24.1949786225266372; 7.78835575489808818; 10.7586857487961058;\n 24.6854987844132197; 13.9889266783498822; 17.7608231677881321;\n 27.4863101958301748; 8.63334537708151; 13.5941912097414512;\n 13.2260260168708097; 31.3766650710703807; 18.9526720076696868;\n 14.0314032315368173; 12.807834225708806; 12.5233417493347;\n 20.3469799530605187; 19.6825161300159905; 14.7618831786326936;\n 34.3446168784105055; 14.372949787259202; 21.4573053352920269;\n 21.0965445810056522; 23.4345755202291386; 14.8926270517011119;\n 14.5962029388981858; 13.3160449309237592; 22.718990622509434;\n 30.5913289186968882; 6.13605093366510346; 12.9969161832653164;\n 27.6045671236346415; 26.2947120504520981; 10.3837082466899;\n 6.29789856614867727; 19.9179299665683978; 23.0504729450445041;\n 29.6802504840401795; 18.5389364101218526; 26.1866389989169512;\n 29.4897237317334593; 25.4911194943350203; 18.0761901209589091;\n 24.5200237809392654; 38.0968922623001305; 22.6861065677917715;\n 34.2502130892676817; 27.9615875110153667; 18.0223599013189144;\n 18.6980019826901689; 18.2672101541450651; 21.8778636636524624;\n 30.0375605276369; 9.91408611351553; 17.2822431119072597;\n 19.1537670175515231; 13.8683272198277763; 12.0528115008079286;\n 30.2168128321327174; 23.5658128966585103; 31.9593800666736279;\n 27.6431867840131886; 9.00369113498522466; 10.6459491096232455;\n 17.612506723563655; 23.941354096378312; 41.4573139679420777;\n 32.7184654368981498; 29.1289053383799548; 17.2569811896428149;\n 13.8128200856228514; 32.0274222366012324; 28.9246665418856;\n 18.3059966853358169; 27.6492920330522622; 8.41661752676145625;\n 24.9548113528369981; 24.8777005517676493; 31.0208268317396296;\n 25.4453634768141832; 8.2424844909615409; 31.304378738480434;\n 11.9805803972240117; 9.57136742791041328; 7.21284524967817653;\n 31.0063188509245471; 29.7391957624394863; ...|]\nval sum_w : float = 38213.713851628665\nval gl_n : float array =\n [|-0.989400934991649939; -0.9445750230732326; -0.865631202387831755;\n -0.755404408355003; -0.617876244402643771; -0.458016777657227425;\n -0.281603550779258915; -0.0950125098376374; 0.0950125098376374;\n 0.281603550779258915; 0.458016777657227425; 0.617876244402643771;\n 0.755404408355003; 0.865631202387831755; 0.9445750230732326;\n 0.989400934991649939|]\nval gl_wt : float array =\n [|0.0271524594117541; 0.0622535239386479; 0.0951585116824928;\n 0.124628971255533905; 0.149595988816576708; 0.169156519395002508;\n 0.182603415044923612; 0.189450610455068502; 0.189450610455068502;\n 0.182603415044923612; 0.169156519395002508; 0.149595988816576708;\n 0.124628971255533905; 0.0951585116824928; 0.0622535239386479;\n 0.0271524594117541|]\nval dist_mod_f : float -> float -> float -> float = \nval chi2_at : float -> float -> float = \nval n_om : int = 100\nval n_ol : int = 100\nval om_min : float = 0.\nval om_max : float = 3.\nval ol_min : float = -1.\nval ol_max : float = 3.\nval chi2_grid : (float, Nx.float64_elt) Nx.t = float64 [100; 100] \n [[1522.01, 1543.98, ..., 3820.7, 3844.02],\n [1484.88, 1506.67, ..., 3779.71, 3803.06],\n ...\n [inf, inf, ..., 877.663, 867.429],\n [inf, inf, ..., 902.83, 890.429]]\nval chi2_min : float = 684.197413463096268\nval delta_chi2 : (float, Nx.float64_elt) Nx.t = float64 [100; 100] \n [[837.808, 859.782, ..., 3136.5, 3159.83],\n [800.684, 822.469, ..., 3095.52, 3118.86],\n ...\n [inf, inf, ..., 193.466, 183.232],\n [inf, inf, ..., 218.632, 206.232]]\n\n\n## Contour plot\n\nReproducing Perlmutter et al. (1999) Figure 7. The contour levels correspond\nto 68%, 90%, 95%, and 99% confidence regions for two parameters\n($\\Delta\\chi^2 = 2.30, 4.61, 5.99, 9.21$). The diagonal solid line marks\n**flat** universes ($\\Omega_M + \\Omega_\\Lambda = 1$). The nearly horizontal\ndashed line separates eternally expanding universes from those that eventually\nrecollapse. The upper-left gray region has no Big Bang.\n\n\n\n```ocaml\n(* Confidence levels for 2 parameters:\n 68% -> delta_chi2 = 2.30, 90% -> 4.61, 95% -> 5.99, 99% -> 9.21 *)\nlet confidence_levels = [| 2.30; 4.61; 5.99; 9.21 |]\n\n(* \"No Big Bang\" boundary: upper-left region where the universe has no\n initial singularity. Approximate as OmegaL > 4*OmegaM*(cosh(...))^3\n for plotting purposes; simplified to a polygon here. *)\nlet no_bb_x = Nx.create f32 [| 5 |] [| 0.0; 0.0; 1.0; 2.0; 3.0 |]\nlet no_bb_y1 = Nx.create f32 [| 5 |] [| 3.0; 1.0; 2.2; 2.8; 3.0 |]\nlet no_bb_y2 = Nx.full f32 [| 5 |] 3.0\n\nlet _fig =\n Hugin.layers [\n (* \"No Big Bang\" shaded region *)\n Hugin.fill_between ~x:no_bb_x ~y1:no_bb_y1 ~y2:no_bb_y2\n ~color:(Hugin.Color.with_alpha 0.15 Hugin.Color.gray) () ;\n (* Filled confidence contours -- blue/teal like Figure 7 *)\n Hugin.contour ~data:(to32 delta_chi2)\n ~x0:om_min ~x1:om_max ~y0:ol_min ~y1:ol_max\n ~levels:(`Values confidence_levels)\n ~filled:true\n ~cmap:(Hugin.Cmap.of_colors [|\n Hugin.Color.with_alpha 0.8 (Hugin.Color.hex \"#1a5276\");\n Hugin.Color.with_alpha 0.6 (Hugin.Color.hex \"#2e86c1\");\n Hugin.Color.with_alpha 0.4 (Hugin.Color.hex \"#85c1e9\");\n Hugin.Color.with_alpha 0.2 (Hugin.Color.hex \"#d4e6f1\");\n Hugin.Color.with_alpha 0.0 Hugin.Color.white;\n |]) () ;\n (* Contour outlines *)\n Hugin.contour ~data:(to32 delta_chi2)\n ~x0:om_min ~x1:om_max ~y0:ol_min ~y1:ol_max\n ~levels:(`Values confidence_levels)\n ~color:(Hugin.Color.hex \"#2e86c1\") ~line_width:1.0 () ;\n (* Flat universe line: OmegaM + OmegaL = 1 *)\n Hugin.line\n ~x:(Nx.create f32 [| 2 |] [| 0.0; 3.0 |])\n ~y:(Nx.create f32 [| 2 |] [| 1.0; -2.0 |])\n ~color:Hugin.Color.black ~line_width:1.5\n ~label:\"Flat\" () ;\n (* No-deceleration line: q0 = 0, i.e. OmegaL = OmegaM / 2 *)\n Hugin.line\n ~x:(Nx.create f32 [| 2 |] [| 0.0; 3.0 |])\n ~y:(Nx.create f32 [| 2 |] [| 0.0; 1.5 |])\n ~color:Hugin.Color.gray ~line_style:`Dashed ~line_width:1.0\n ~label:\"Accelerating/decelerating\" () ;\n (* Lambda = 0 line *)\n Hugin.hline ~y:0.0 ~color:Hugin.Color.gray ~line_style:`Dotted\n ~line_width:0.5 () ;\n ]\n |> Hugin.xlim 0.0 3.0\n |> Hugin.ylim (-1.0) 3.0\n |> Hugin.xlabel \"Omega_M\"\n |> Hugin.ylabel \"Omega_Lambda\"\n |> Hugin.title \"Confidence Contours in the Omega_M - Omega_Lambda Plane\"\n |> Hugin.legend ~loc:Hugin.Upper_right\n```\n\n\nval confidence_levels : float array = [|2.3; 4.61; 5.99; 9.21|]\nval no_bb_x : (float, Nx.float32_elt) Nx.t = float32 [5] [0, 0, ..., 2, 3]\nval no_bb_y1 : (float, Nx.float32_elt) Nx.t = float32 [5] [3, 1, ..., 2.8, 3]\nval no_bb_y2 : (float, Nx.float32_elt) Nx.t = float32 [5] [3, 3, ..., 3, 3]\nval _fig : Hugin.t =\n \n\n\n\n\n## Conclusion\n\nWe have reproduced the central result of Perlmutter et al. (1999) using the\nmodern Pantheon+ dataset and Umbra's cosmology module:\n\n1. The **Hubble diagram** shows that distant SNe Ia are fainter than predicted\n by decelerating models, confirming cosmic acceleration.\n2. **Residuals** relative to an empty universe clearly show the acceleration\n signal at $z > 0.2$.\n3. **Confidence contours** in the $\\Omega_M$--$\\Omega_\\Lambda$ plane strongly\n exclude $\\Omega_\\Lambda = 0$ and are consistent with a flat universe with\n $\\Omega_M \\approx 0.3$, $\\Omega_\\Lambda \\approx 0.7$.\n\nThe analysis required only Umbra's `Cosmo.lcdm`, `Cosmo.flat_lcdm`, and\n`Cosmo.distance_modulus` functions -- the entire theoretical framework for\nSN Ia cosmology in a few lines of OCaml.\n\n### References\n\n- Perlmutter, S. et al. 1999, ApJ, 517, 565 (arXiv:astro-ph/9812133)\n- Riess, A.G. et al. 1998, AJ, 116, 1009 (arXiv:astro-ph/9805201)\n- Scolnic, D.M. et al. 2022, ApJ, 938, 113 (arXiv:2112.03863)\n- Brout, D. et al. 2022, ApJ, 938, 110 (arXiv:2202.04077)\n"
},
{
"path": "dev/umbra/test/dune",
"content": "(test\n (name test_umbra)\n (libraries umbra umbra.fits nx nx.io talon windtrap))\n"
},
{
"path": "dev/umbra/test/test_umbra.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nopen Umbra\n\nlet eps = 1e-6\nlet f64 = Nx.float64\nlet v x = Nx.item [] x\n\n(* Unit tests *)\n\nlet test_length_conversion () =\n let d = Unit.Length.kpc 10.0 in\n let mpc = v (Unit.Length.in_mpc d) in\n is_true ~msg:\"10 kpc = 0.01 Mpc\" (Float.abs (mpc -. 0.01) < eps);\n let m = v (Unit.Length.in_m d) in\n let back = v (Unit.Length.in_kpc (Unit.Length.m m)) in\n is_true ~msg:\"kpc -> m -> kpc roundtrip\" (Float.abs (back -. 10.0) < eps)\n\nlet test_length_arithmetic () =\n let open Unit in\n let d = Length.kpc 10.0 + Length.pc 500.0 in\n let kpc = v (Length.in_kpc d) in\n is_true ~msg:\"10 kpc + 500 pc = 10.5 kpc\" (Float.abs (kpc -. 10.5) < eps)\n\nlet test_mass_conversion () =\n let m = Unit.Mass.solar_mass 1.0 in\n let kg = v (Unit.Mass.in_kg m) in\n is_true ~msg:\"1 Msun ~ 1.988e30 kg\"\n (Float.abs (kg -. 1.9884e30) /. 1.9884e30 < 1e-4)\n\nlet test_velocity_cross_dim () =\n let d = Unit.Length.km 100.0 in\n let t = Unit.Time.s 10.0 in\n let vel = Unit.length_per_time d t in\n let km_s = v (Unit.Velocity.in_km_s vel) in\n is_true ~msg:\"100 km / 10 s = 10 km/s\" (Float.abs (km_s -. 10.0) < eps)\n\nlet test_angle_trig () =\n let a = Unit.Angle.deg 90.0 in\n is_true ~msg:\"sin(90°) = 1\"\n (Float.abs (Nx.item [] (Unit.Angle.sin a) -. 1.0) < eps);\n is_true ~msg:\"cos(90°) = 0\" (Float.abs (Nx.item [] (Unit.Angle.cos a)) < eps)\n\nlet test_wavelength_frequency () =\n let lam = Unit.Length.nm 500.0 in\n let nu = Unit.wavelength_to_frequency lam in\n let lam2 = Unit.frequency_to_wavelength nu in\n let nm2 = v (Unit.Length.in_nm lam2) in\n is_true ~msg:\"wavelength -> freq -> wavelength roundtrip\"\n (Float.abs (nm2 -. 500.0) < eps)\n\nlet test_phantom_type_safety () =\n (* This is a compile-time test: the following should NOT typecheck: let _ =\n Unit.(Length.m 1.0 + Mass.kg 1.0) The fact that this module compiles proves\n type safety. *)\n let _d = Unit.(Length.m 1.0 + Length.km 1.0) in\n let _m = Unit.(Mass.kg 1.0 + Mass.g 500.0) in\n ()\n\n(* Const tests *)\n\nlet test_const_c () =\n let c_km_s = v (Unit.Velocity.in_km_s Const.c) in\n is_true ~msg:\"c ~ 299792 km/s\" (Float.abs (c_km_s -. 299792.458) < 1.0)\n\n(* Coord tests *)\n\nlet deg_eps = 1e-6\n\nlet test_coord_roundtrip () =\n let ra =\n Unit.Angle.of_deg (Nx.create f64 [| 4 |] [| 180.0; 0.0; 90.0; 266.405 |])\n in\n let dec =\n Unit.Angle.of_deg (Nx.create f64 [| 4 |] [| 45.0; -30.0; 0.0; -28.936 |])\n in\n let c = Coord.of_radec ~ra ~dec in\n let gal = Coord.galactic c in\n let back = Coord.icrs gal in\n let ra' = Unit.Angle.in_deg (Coord.ra back) in\n let dec' = Unit.Angle.in_deg (Coord.dec back) in\n let ra_orig = Unit.Angle.in_deg ra in\n let dec_orig = Unit.Angle.in_deg dec in\n for i = 0 to 3 do\n is_true\n ~msg:(Printf.sprintf \"RA roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] ra_orig -. Nx.item [ i ] ra') < deg_eps);\n is_true\n ~msg:(Printf.sprintf \"Dec roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] dec_orig -. Nx.item [ i ] dec') < deg_eps)\n done\n\nlet test_separation_poles () =\n let c1 =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 0.0 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 90.0 |]))\n in\n let c2 =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 0.0 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| -90.0 |]))\n in\n let sep = Coord.separation c1 c2 in\n is_true ~msg:\"Pole separation = 180\"\n (Float.abs (Nx.item [ 0 ] (Unit.Angle.in_deg sep) -. 180.0) < deg_eps)\n\n(* Cosmo tests *)\n\nlet test_cosmo_distances () =\n let z = Nx.scalar f64 0.1 in\n let dc = v (Unit.Length.in_mpc (Cosmo.comoving_distance z)) in\n is_true\n ~msg:(Printf.sprintf \"comoving(0.1) ~ 421 Mpc, got %.1f\" dc)\n (Float.abs (dc -. 421.0) < 5.0);\n let dl = v (Unit.Length.in_mpc (Cosmo.luminosity_distance z)) in\n is_true\n ~msg:(Printf.sprintf \"luminosity(0.1) ~ 463 Mpc, got %.1f\" dl)\n (Float.abs (dl -. 463.0) < 5.0)\n\nlet test_cosmo_lookback () =\n let z = Nx.scalar f64 1.0 in\n let t = v (Unit.Time.in_gyr (Cosmo.lookback_time z)) in\n is_true\n ~msg:(Printf.sprintf \"lookback(1.0) ~ 7.7 Gyr, got %.1f\" t)\n (Float.abs (t -. 7.7) < 0.3)\n\nlet test_cosmo_angular_scale () =\n let phys = Unit.Length.kpc 1.0 in\n let z = Nx.scalar f64 0.022 in\n let ang = Cosmo.angular_size ~z phys in\n let arcsec = v (Unit.Angle.in_arcsec ang) in\n is_true\n ~msg:(Printf.sprintf \"1 kpc at z=0.022 ~ 2.3 arcsec, got %.2f\" arcsec)\n (Float.abs (arcsec -. 2.3) < 0.2)\n\n(* Cosmo: high-z regression tests. These catch quadrature under-resolution at\n large z. *)\n\nlet test_cosmo_age_planck18 () =\n let p = Cosmo.planck18 in\n let t = v (Unit.Time.in_gyr (Cosmo.age ~p (Nx.scalar f64 0.0))) in\n is_true\n ~msg:(Printf.sprintf \"age(Planck18, z=0) ~ 13.8 Gyr, got %.1f\" t)\n (Float.abs (t -. 13.8) < 0.3)\n\nlet test_cosmo_age_at_z1 () =\n let p = Cosmo.planck18 in\n let age_0 = v (Unit.Time.in_gyr (Cosmo.age ~p (Nx.scalar f64 0.0))) in\n let age_1 = v (Unit.Time.in_gyr (Cosmo.age ~p (Nx.scalar f64 1.0))) in\n let lb_1 =\n v (Unit.Time.in_gyr (Cosmo.lookback_time ~p (Nx.scalar f64 1.0)))\n in\n is_true\n ~msg:\n (Printf.sprintf\n \"age(z=0) - age(z=1) = lookback(z=1): %.2f - %.2f = %.2f vs %.2f\" age_0\n age_1 (age_0 -. age_1) lb_1)\n (Float.abs (age_0 -. age_1 -. lb_1) < 0.05)\n\nlet test_cosmo_comoving_cmb () =\n let p = Cosmo.planck18 in\n let dc =\n v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p (Nx.scalar f64 1089.0)))\n in\n is_true\n ~msg:(Printf.sprintf \"comoving(z=1089) ~ 14000 Mpc, got %.0f\" dc)\n (Float.abs (dc -. 14000.0) < 500.0)\n\nlet test_cosmo_comoving_high_z () =\n let p = Cosmo.planck18 in\n let dc_2 =\n v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p (Nx.scalar f64 2.0)))\n in\n let dc_5 =\n v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p (Nx.scalar f64 5.0)))\n in\n let dc_10 =\n v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p (Nx.scalar f64 10.0)))\n in\n is_true ~msg:\"comoving distances monotonically increase\"\n (dc_2 < dc_5 && dc_5 < dc_10);\n is_true\n ~msg:(Printf.sprintf \"comoving(z=10) ~ 9700 Mpc, got %.0f\" dc_10)\n (Float.abs (dc_10 -. 9700.0) < 300.0)\n\nlet test_cosmo_lookback_high_z () =\n let p = Cosmo.planck18 in\n let lb_5 =\n v (Unit.Time.in_gyr (Cosmo.lookback_time ~p (Nx.scalar f64 5.0)))\n in\n is_true\n ~msg:(Printf.sprintf \"lookback(z=5) ~ 12.5 Gyr, got %.1f\" lb_5)\n (Float.abs (lb_5 -. 12.5) < 0.3)\n\n(* FITS tests *)\n\nlet test_fits_image_roundtrip () =\n let path = \"_test_image.fits\" in\n Fun.protect\n ~finally:(fun () -> if Sys.file_exists path then Sys.remove path)\n (fun () ->\n let data =\n Nx.create Nx.float32 [| 2; 3 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0 |]\n in\n Umbra_fits.write_image path data;\n let packed = Umbra_fits.read_image ~hdu:0 path in\n let result = Nx_io.to_typed Nx.float32 packed in\n is_true ~msg:\"Image shape\" (Nx.shape result = [| 2; 3 |]);\n for i = 0 to 5 do\n let row = i / 3 and col = i mod 3 in\n is_true\n ~msg:(Printf.sprintf \"Image value [%d,%d]\" row col)\n (Float.abs (Nx.item [ row; col ] data -. Nx.item [ row; col ] result)\n < 1e-6)\n done)\n\nlet test_fits_table_roundtrip () =\n let path = \"_test_table.fits\" in\n Fun.protect\n ~finally:(fun () -> if Sys.file_exists path then Sys.remove path)\n (fun () ->\n let df =\n Talon.create\n [\n (\"ra\", Talon.Col.float64 [| 10.0; 20.0; 30.0 |]);\n (\"dec\", Talon.Col.float64 [| -10.0; 0.0; 10.0 |]);\n ]\n in\n Umbra_fits.write_table path df;\n let df2 = Umbra_fits.read_table ~hdu:1 path in\n is_true ~msg:\"Table rows\" (Talon.num_rows df2 = 3);\n match Talon.to_array Nx.float64 df2 \"ra\" with\n | Some arr -> is_true ~msg:\"ra[0]\" (Float.abs (arr.(0) -. 10.0) < 1e-10)\n | None -> fail \"ra column missing\")\n\n(* Coord cross-matching tests *)\n\nlet test_match_nearest_self () =\n let ra = Unit.Angle.of_deg (Nx.create f64 [| 3 |] [| 10.0; 20.0; 30.0 |]) in\n let dec = Unit.Angle.of_deg (Nx.create f64 [| 3 |] [| -10.0; 0.0; 10.0 |]) in\n let c = Coord.of_radec ~ra ~dec in\n let { Coord.indices; separations } = Coord.nearest c c in\n for i = 0 to 2 do\n is_true\n ~msg:(Printf.sprintf \"Self-match index[%d]\" i)\n (Int32.to_int (Nx.item [ i ] indices) = i);\n is_true\n ~msg:(Printf.sprintf \"Self-match separation[%d]\" i)\n (Nx.item [ i ] (Unit.Angle.in_rad separations) < 1e-10)\n done\n\n(* Time tests *)\n\nlet test_time_jd_mjd () =\n let t = Time.unsafe_of_jd 2451545.0 in\n is_true ~msg:\"J2000.0 JD\" (Float.abs (Time.to_jd t -. 2451545.0) < 1e-10);\n is_true ~msg:\"J2000.0 MJD\" (Float.abs (Time.to_mjd t -. 51544.5) < 1e-10);\n let t2 = Time.unsafe_of_mjd 51544.5 in\n is_true ~msg:\"MJD roundtrip\" (Float.abs (Time.to_jd t2 -. 2451545.0) < 1e-10)\n\nlet test_time_iso () =\n let t = Time.of_iso \"2000-01-01T12:00:00\" in\n is_true ~msg:\"J2000.0 from ISO\" (Float.abs (Time.to_jd t -. 2451545.0) < 1e-6);\n let s = Time.to_iso t in\n is_true ~msg:\"ISO roundtrip\" (s = \"2000-01-01T12:00:00Z\")\n\nlet test_time_utc_tai_tt () =\n let utc = Time.unsafe_of_jd 2451545.0 in\n let tai = Time.utc_to_tai utc in\n let dt_s = (Time.to_jd tai -. Time.to_jd utc) *. 86400.0 in\n is_true\n ~msg:(Printf.sprintf \"TAI-UTC at J2000 = 32s, got %.1f\" dt_s)\n (Float.abs (dt_s -. 32.0) < 0.1);\n let tt = Time.tai_to_tt tai in\n let dt_tt = (Time.to_jd tt -. Time.to_jd tai) *. 86400.0 in\n is_true\n ~msg:(Printf.sprintf \"TT-TAI = 32.184s, got %.6f\" dt_tt)\n (Float.abs (dt_tt -. 32.184) < 1e-3);\n let tai' = Time.tt_to_tai tt in\n is_true ~msg:\"TT->TAI roundtrip\"\n (Float.abs (Time.to_jd tai' -. Time.to_jd tai) < 1e-12);\n let utc' = Time.tai_to_utc tai in\n is_true ~msg:\"TAI->UTC roundtrip\"\n (Float.abs (Time.to_jd utc' -. Time.to_jd utc) < 1e-10)\n\nlet test_time_tdb () =\n let tt = Time.unsafe_of_jd 2451545.0 in\n let tdb = Time.tt_to_tdb tt in\n let dt_ms = (Time.to_jd tdb -. Time.to_jd tt) *. 86400.0 *. 1000.0 in\n is_true\n ~msg:(Printf.sprintf \"TDB-TT < 2ms, got %.3f ms\" dt_ms)\n (Float.abs dt_ms < 2.0);\n let tt' = Time.tdb_to_tt tdb in\n is_true ~msg:\"TDB->TT roundtrip\"\n (Float.abs (Time.to_jd tt' -. Time.to_jd tt) < 1e-10)\n\nlet test_time_unix () =\n let t = Time.of_unix 0.0 in\n is_true ~msg:\"Unix epoch JD\" (Float.abs (Time.to_jd t -. 2440587.5) < 1e-10);\n let u = Time.to_unix t in\n is_true ~msg:\"Unix roundtrip\" (Float.abs u < 1e-6)\n\nlet test_time_diff_add () =\n let t1 = Time.unsafe_of_jd 2451545.0 in\n let t2 = Time.unsafe_of_jd 2451546.0 in\n let dt = Time.diff t2 t1 in\n is_true ~msg:\"diff = 1 day\"\n (Float.abs (v (Unit.Time.in_day dt) -. 1.0) < 1e-10);\n let t3 = Time.add t1 (Unit.Time.day 1.0) in\n is_true ~msg:\"add 1 day\" (Float.abs (Time.to_jd t3 -. 2451546.0) < 1e-10)\n\n(* Cosmo preset tests *)\n\nlet test_cosmo_planck18 () =\n let z = Nx.scalar f64 0.5 in\n let dc =\n v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p:Cosmo.planck18 z))\n in\n is_true\n ~msg:(Printf.sprintf \"Planck18 comoving(0.5) ~ 1960 Mpc, got %.0f\" dc)\n (Float.abs (dc -. 1960.0) < 30.0)\n\nlet test_cosmo_hubble () =\n let z = Nx.scalar f64 0.0 in\n let h0 = Nx.item [] (Cosmo.hubble z) in\n is_true\n ~msg:(Printf.sprintf \"H(0) = H0 = 70, got %.1f\" h0)\n (Float.abs (h0 -. 70.0) < 1e-6)\n\n(* Coord FK5/Supergalactic tests *)\n\nlet test_coord_ecliptic_roundtrip () =\n let ra = Unit.Angle.of_deg (Nx.create f64 [| 2 |] [| 180.0; 45.0 |]) in\n let dec = Unit.Angle.of_deg (Nx.create f64 [| 2 |] [| 45.0; -30.0 |]) in\n let c = Coord.of_radec ~ra ~dec in\n let ecl = Coord.ecliptic_j2000 c in\n let back = Coord.icrs ecl in\n let ra' = Unit.Angle.in_deg (Coord.ra back) in\n let dec' = Unit.Angle.in_deg (Coord.dec back) in\n let ra_orig = Unit.Angle.in_deg ra in\n let dec_orig = Unit.Angle.in_deg dec in\n for i = 0 to 1 do\n is_true\n ~msg:(Printf.sprintf \"Ecliptic RA roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] ra_orig -. Nx.item [ i ] ra') < deg_eps);\n is_true\n ~msg:(Printf.sprintf \"Ecliptic Dec roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] dec_orig -. Nx.item [ i ] dec') < deg_eps)\n done\n\nlet test_coord_supergalactic_roundtrip () =\n let ra = Unit.Angle.of_deg (Nx.create f64 [| 2 |] [| 180.0; 45.0 |]) in\n let dec = Unit.Angle.of_deg (Nx.create f64 [| 2 |] [| 45.0; -30.0 |]) in\n let c = Coord.of_radec ~ra ~dec in\n let sg = Coord.supergalactic c in\n let back = Coord.icrs sg in\n let ra' = Unit.Angle.in_deg (Coord.ra back) in\n let dec' = Unit.Angle.in_deg (Coord.dec back) in\n let ra_orig = Unit.Angle.in_deg ra in\n let dec_orig = Unit.Angle.in_deg dec in\n for i = 0 to 1 do\n is_true\n ~msg:(Printf.sprintf \"Supergalactic RA roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] ra_orig -. Nx.item [ i ] ra') < 1e-4);\n is_true\n ~msg:(Printf.sprintf \"Supergalactic Dec roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] dec_orig -. Nx.item [ i ] dec') < 1e-4)\n done\n\n(* Unit energy-wavelength-frequency tests *)\n\nlet test_energy_wavelength_frequency () =\n let e = Unit.Energy.ev 2.0 in\n let nu = Unit.energy_to_frequency e in\n let e2 = Unit.frequency_to_energy nu in\n is_true ~msg:\"energy->freq->energy roundtrip\"\n (Float.abs (v (Unit.Energy.in_ev e2) -. 2.0) < 1e-6);\n let lam = Unit.energy_to_wavelength e in\n let nu2 = Unit.wavelength_to_frequency lam in\n let e3 = Unit.frequency_to_energy nu2 in\n is_true ~msg:\"energy->wavelength->freq->energy roundtrip\"\n (Float.abs (v (Unit.Energy.in_ev e3) -. 2.0) < 1e-6)\n\n(* Spectrum tests *)\n\nlet test_spectrum_blackbody_wien () =\n (* Wien's displacement law: λ_max * T = 2.898e-3 m·K *)\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 5778.0) in\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 3e-6 1000) in\n let spec = Spectrum.blackbody ~temperature:temp ~wavelength:wave in\n let vals = Spectrum.values spec in\n (* Find index of max value *)\n let peak_idx = ref 0 in\n let peak_val = ref (Nx.item [ 0 ] vals) in\n for i = 1 to 999 do\n let v = Nx.item [ i ] vals in\n if v > !peak_val then begin\n peak_val := v;\n peak_idx := i\n end\n done;\n let wave_m = Unit.Length.in_m (Spectrum.wavelength spec) in\n let peak_lam = Nx.item [ !peak_idx ] wave_m in\n let wien = peak_lam *. 5778.0 in\n is_true\n ~msg:(Printf.sprintf \"Wien's law: λ_max*T ~ 2.898e-3, got %.4e\" wien)\n (Float.abs (wien -. 2.898e-3) /. 2.898e-3 < 0.01)\n\nlet test_spectrum_redshift () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-6 100) in\n let values = Nx.ones f64 [| 100 |] in\n let spec =\n Spectrum.create ~wavelength:wave ~values |> Spectrum.as_flux_density\n in\n let z = Nx.scalar f64 1.0 in\n let shifted = Spectrum.redshift ~z spec in\n (* Wavelengths should double at z=1 *)\n let orig_wave = Unit.Length.in_m (Spectrum.wavelength spec) in\n let shifted_wave = Unit.Length.in_m (Spectrum.wavelength shifted) in\n let ratio = Nx.item [ 50 ] shifted_wave /. Nx.item [ 50 ] orig_wave in\n is_true ~msg:\"Redshift z=1 doubles wavelength\"\n (Float.abs (ratio -. 2.0) < 1e-10);\n (* Values should halve at z=1 *)\n let val_ratio =\n Nx.item [ 50 ] (Spectrum.values shifted) /. Nx.item [ 50 ] values\n in\n is_true ~msg:\"Redshift z=1 halves values\"\n (Float.abs (val_ratio -. 0.5) < 1e-10)\n\nlet test_spectrum_scale () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-6 10) in\n let values = Nx.ones f64 [| 10 |] in\n let spec = Spectrum.create ~wavelength:wave ~values in\n let scaled = Spectrum.scale (Nx.scalar f64 3.0) spec in\n is_true ~msg:\"Scale by 3\"\n (Float.abs (Nx.item [ 0 ] (Spectrum.values scaled) -. 3.0) < 1e-10)\n\n(* Extinction tests *)\n\nlet test_extinction_ccm89_v_band () =\n (* At V-band (550nm), A_λ/A_V should be ~1.0 for R_V=3.1 *)\n let rv = Nx.scalar f64 3.1 in\n let wave_v = Unit.Length.of_m (Nx.create f64 [| 1 |] [| 5.5e-7 |]) in\n let alav = Extinction.curve (Extinction.ccm89 ~rv) ~wavelength:wave_v in\n let val_v = Nx.item [ 0 ] alav in\n is_true\n ~msg:(Printf.sprintf \"CCM89 A_V/A_V ~ 1.0 at 550nm, got %.3f\" val_v)\n (Float.abs (val_v -. 1.0) < 0.1)\n\nlet test_extinction_apply_unredden () =\n (* apply then unredden should recover original spectrum *)\n let rv = Nx.scalar f64 3.1 in\n let law = Extinction.ccm89 ~rv in\n let wave = Unit.Length.of_m (Nx.linspace f64 3e-7 1e-6 50) in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 6000.0) in\n let spec = Spectrum.blackbody ~temperature:temp ~wavelength:wave in\n let av = Nx.scalar f64 1.0 in\n let reddened = Extinction.apply law ~av spec in\n let recovered = Extinction.unredden law ~av reddened in\n (* Compare values *)\n let orig_val = Nx.item [ 25 ] (Spectrum.values spec) in\n let rec_val = Nx.item [ 25 ] (Spectrum.values recovered) in\n is_true ~msg:\"apply + unredden roundtrip\"\n (Float.abs (rec_val -. orig_val) /. orig_val < 1e-10)\n\nlet test_extinction_ccm89_monotonic () =\n (* Extinction should increase toward blue wavelengths (for optical) *)\n let rv = Nx.scalar f64 3.1 in\n let wave =\n Unit.Length.of_m (Nx.create f64 [| 3 |] [| 4e-7; 5.5e-7; 8e-7 |])\n in\n let alav = Extinction.curve (Extinction.ccm89 ~rv) ~wavelength:wave in\n let blue = Nx.item [ 0 ] alav in\n let green = Nx.item [ 1 ] alav in\n let red = Nx.item [ 2 ] alav in\n is_true ~msg:\"CCM89: A_blue > A_green\" (blue > green);\n is_true ~msg:\"CCM89: A_green > A_red\" (green > red)\n\n(* Photometry tests *)\n\nlet test_photometry_ab_mag_flat () =\n (* A flat f_nu spectrum at 3631 Jy should give m_AB = 0 in any band. f_nu =\n 3631e-26 W/m²/Hz, so f_lambda = f_nu * c / lambda² *)\n let n = 100 in\n let bp =\n Photometry.tophat ~lo:(Unit.Length.m 4e-7) ~hi:(Unit.Length.m 7e-7) ~n\n in\n let wave_m = Unit.Length.to_tensor (Photometry.wavelength bp) in\n let c = 299_792_458.0 in\n let ab_zp = 3631.0e-26 in\n (* f_lambda = f_nu * c / lambda^2 *)\n let f_lambda =\n Nx.div\n (Nx.mul_s (Nx.recip (Nx.square wave_m)) (ab_zp *. c))\n (Nx.scalar f64 1.0)\n in\n let spec =\n Spectrum.create ~wavelength:(Photometry.wavelength bp) ~values:f_lambda\n |> Spectrum.as_flux_density\n in\n let mag = Nx.item [] (Photometry.ab_mag bp spec) in\n is_true\n ~msg:(Printf.sprintf \"Flat f_nu=3631Jy gives m_AB ~ 0, got %.3f\" mag)\n (Float.abs mag < 0.05)\n\nlet test_photometry_color_same_band () =\n (* Color between same band should be 0 *)\n let bp =\n Photometry.tophat ~lo:(Unit.Length.m 4e-7) ~hi:(Unit.Length.m 5.5e-7) ~n:50\n in\n let spec =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 5000.0))\n ~wavelength:(Photometry.wavelength bp)\n |> Spectrum.as_flux_density\n in\n let col = Nx.item [] (Photometry.color bp bp spec) in\n is_true ~msg:\"Same-band color = 0\" (Float.abs col < 1e-10)\n\nlet test_photometry_blue_star_color () =\n (* A hot star should be brighter (lower mag) in blue than red *)\n let n = 100 in\n let bp_b =\n Photometry.tophat ~lo:(Unit.Length.m 4e-7) ~hi:(Unit.Length.m 5e-7) ~n\n in\n let bp_r =\n Photometry.tophat ~lo:(Unit.Length.m 6e-7) ~hi:(Unit.Length.m 7e-7) ~n\n in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 20000.0) in\n let spec_b =\n Spectrum.blackbody ~temperature:temp\n ~wavelength:(Photometry.wavelength bp_b)\n |> Spectrum.as_flux_density\n in\n let spec_r =\n Spectrum.blackbody ~temperature:temp\n ~wavelength:(Photometry.wavelength bp_r)\n |> Spectrum.as_flux_density\n in\n let mag_b = Nx.item [] (Photometry.ab_mag bp_b spec_b) in\n let mag_r = Nx.item [] (Photometry.ab_mag bp_r spec_r) in\n is_true ~msg:\"Hot star: blue mag < red mag (brighter in blue)\" (mag_b < mag_r)\n\n(* Cosmo: extended models *)\n\nlet test_cosmo_flat_lcdm_same_as_default () =\n let p = Cosmo.flat_lcdm ~h0:70.0 ~omega_m:0.3 in\n let z = Nx.scalar f64 0.5 in\n let dc_default = v (Unit.Length.in_mpc (Cosmo.comoving_distance z)) in\n let dc_flat = v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p z)) in\n is_true ~msg:\"flat_lcdm(70,0.3) = default\"\n (Float.abs (dc_default -. dc_flat) < 1e-6)\n\nlet test_cosmo_nonflat_lcdm () =\n (* Open universe: omega_m=0.3, omega_l=0.5 → omega_k=0.2. Result should differ\n from flat LCDM. *)\n let p_flat = Cosmo.flat_lcdm ~h0:70.0 ~omega_m:0.3 in\n let p_open = Cosmo.lcdm ~h0:70.0 ~omega_m:0.3 ~omega_l:0.5 in\n let z = Nx.scalar f64 1.0 in\n let dl_flat =\n v (Unit.Length.in_mpc (Cosmo.luminosity_distance ~p:p_flat z))\n in\n let dl_open =\n v (Unit.Length.in_mpc (Cosmo.luminosity_distance ~p:p_open z))\n in\n is_true\n ~msg:\n (Printf.sprintf \"Non-flat LCDM differs from flat: %.0f vs %.0f\" dl_open\n dl_flat)\n (Float.abs (dl_open -. dl_flat) > 10.0)\n\nlet test_cosmo_wcdm () =\n (* w0 = -1 should be identical to ΛCDM *)\n let p_lcdm = Cosmo.flat_lcdm ~h0:70.0 ~omega_m:0.3 in\n let p_wcdm = Cosmo.wcdm ~h0:70.0 ~omega_m:0.3 ~w0:(-1.0) () in\n let z = Nx.scalar f64 0.5 in\n let dc_lcdm = v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p:p_lcdm z)) in\n let dc_wcdm = v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p:p_wcdm z)) in\n is_true\n ~msg:(Printf.sprintf \"wCDM(w0=-1) = LCDM: %.1f vs %.1f\" dc_wcdm dc_lcdm)\n (Float.abs (dc_wcdm -. dc_lcdm) < 1.0)\n\nlet test_cosmo_w0wacdm () =\n (* w0=-1, wa=0 should reduce to ΛCDM *)\n let p_lcdm = Cosmo.flat_lcdm ~h0:70.0 ~omega_m:0.3 in\n let p_cpl = Cosmo.w0wacdm ~h0:70.0 ~omega_m:0.3 ~w0:(-1.0) ~wa:0.0 () in\n let z = Nx.scalar f64 1.0 in\n let dl_lcdm =\n v (Unit.Length.in_mpc (Cosmo.luminosity_distance ~p:p_lcdm z))\n in\n let dl_cpl = v (Unit.Length.in_mpc (Cosmo.luminosity_distance ~p:p_cpl z)) in\n is_true\n ~msg:(Printf.sprintf \"w0waCDM(-1,0) = LCDM: %.1f vs %.1f\" dl_cpl dl_lcdm)\n (Float.abs (dl_cpl -. dl_lcdm) < 1.0)\n\nlet test_cosmo_e_of () =\n (* E(z=0) = 1 for any cosmology *)\n let p = Cosmo.planck18 in\n let z = Nx.scalar f64 0.0 in\n let e0 = v (Cosmo.e_of p z) in\n is_true\n ~msg:(Printf.sprintf \"E(z=0) = 1, got %.6f\" e0)\n (Float.abs (e0 -. 1.0) < 1e-6)\n\nlet test_cosmo_z_at_value () =\n (* Roundtrip: compute dl at z=0.5, then find z back *)\n let p = Cosmo.default in\n let z0 = 0.5 in\n let dl = Cosmo.luminosity_distance ~p (Nx.scalar f64 z0) in\n let z_found =\n v\n (Cosmo.z_at_value ~p\n (fun ~p z -> Unit.Length.to_tensor (Cosmo.luminosity_distance ~p z))\n (Unit.Length.to_tensor dl))\n in\n is_true\n ~msg:(Printf.sprintf \"z_at_value roundtrip: expected 0.5, got %.6f\" z_found)\n (Float.abs (z_found -. z0) < 1e-6)\n\n(* AltAz tests *)\n\nlet test_altaz_zenith () =\n (* A star at the observer's zenith should have alt ~ 90° *)\n let obs =\n Altaz.make_observer ~lat:(Unit.Angle.deg 0.0) ~lon:(Unit.Angle.deg 0.0) ()\n in\n (* Use the vernal equinox time: RA=0, Dec=0 should be near zenith at sidereal\n midnight from lon=0, lat=0. At J2000.0 the ERA is ~280.46°, so RA ~ 280.46°\n should be near transit. Instead, test roundtrip. *)\n let t = Time.of_iso \"2024-01-01T00:00:00\" in\n let ra = Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 180.0 |]) in\n let dec = Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 45.0 |]) in\n let c = Coord.of_radec ~ra ~dec in\n let hz = Altaz.of_coord ~obstime:t ~observer:obs c in\n let back = Altaz.to_coord ~obstime:t ~observer:obs hz in\n let ra' = Nx.item [ 0 ] (Unit.Angle.in_deg (Coord.ra back)) in\n let dec' = Nx.item [ 0 ] (Unit.Angle.in_deg (Coord.dec back)) in\n is_true\n ~msg:(Printf.sprintf \"AltAz RA roundtrip: 180 vs %.4f\" ra')\n (Float.abs (ra' -. 180.0) < 0.1);\n is_true\n ~msg:(Printf.sprintf \"AltAz Dec roundtrip: 45 vs %.4f\" dec')\n (Float.abs (dec' -. 45.0) < 0.1)\n\nlet test_altaz_north_pole () =\n (* Polaris (dec ~ 90) should always be near alt = observer lat *)\n let obs =\n Altaz.make_observer ~lat:(Unit.Angle.deg 45.0) ~lon:(Unit.Angle.deg 0.0) ()\n in\n let t = Time.of_iso \"2024-06-15T12:00:00\" in\n let ra = Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 37.95 |]) in\n let dec = Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 89.264 |]) in\n let c = Coord.of_radec ~ra ~dec in\n let hz = Altaz.of_coord ~obstime:t ~observer:obs c in\n let alt_deg = Nx.item [ 0 ] (Unit.Angle.in_deg (Altaz.alt hz)) in\n is_true\n ~msg:(Printf.sprintf \"Polaris alt ~ 45° from lat=45°, got %.1f\" alt_deg)\n (Float.abs (alt_deg -. 45.0) < 2.0)\n\n(* Galactocentric tests *)\n\nlet test_galactocentric_gc_position () =\n (* A point at l=0, b=0, d=galcen_distance should map to near (0, 0, z_sun) in\n Galactocentric. *)\n let c =\n Coord.of_galactic\n ~l:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 0.0 |]))\n ~b:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 0.0 |]))\n in\n let gc =\n Galactocentric.of_coord\n ~distance:(Unit.Length.of_kpc (Nx.create f64 [| 1 |] [| 8.122 |]))\n c\n in\n let xv = Nx.item [ 0 ] (Unit.Length.in_kpc (Galactocentric.x gc)) in\n let yv = Nx.item [ 0 ] (Unit.Length.in_kpc (Galactocentric.y gc)) in\n let zv = Nx.item [ 0 ] (Unit.Length.in_kpc (Galactocentric.z gc)) in\n is_true\n ~msg:(Printf.sprintf \"GC x ~ 0 kpc, got %.6f\" xv)\n (Float.abs xv < 1e-10);\n is_true\n ~msg:(Printf.sprintf \"GC y ~ 0 kpc, got %.6f\" yv)\n (Float.abs yv < 1e-10);\n is_true\n ~msg:(Printf.sprintf \"GC z ~ z_sun=0.0208 kpc, got %.4f\" zv)\n (Float.abs (zv -. 0.0208) < 1e-10)\n\nlet test_galactocentric_roundtrip () =\n let ra = Unit.Angle.of_deg (Nx.create f64 [| 2 |] [| 180.0; 45.0 |]) in\n let dec = Unit.Angle.of_deg (Nx.create f64 [| 2 |] [| 30.0; -15.0 |]) in\n let c = Coord.of_radec ~ra ~dec in\n let d = Unit.Length.of_kpc (Nx.create f64 [| 2 |] [| 5.0; 12.0 |]) in\n let gc = Galactocentric.of_coord ~distance:d c in\n let c', d' = Galactocentric.to_coord gc in\n let ra' = Unit.Angle.in_deg (Coord.ra c') in\n let dec' = Unit.Angle.in_deg (Coord.dec c') in\n let d_kpc' = Unit.Length.in_kpc d' in\n let ra_orig = Unit.Angle.in_deg ra in\n let dec_orig = Unit.Angle.in_deg dec in\n let d_orig = Unit.Length.in_kpc d in\n for i = 0 to 1 do\n is_true\n ~msg:(Printf.sprintf \"Galactocentric RA roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] ra' -. Nx.item [ i ] ra_orig) < 0.01);\n is_true\n ~msg:(Printf.sprintf \"Galactocentric Dec roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] dec' -. Nx.item [ i ] dec_orig) < 0.01);\n is_true\n ~msg:(Printf.sprintf \"Galactocentric distance roundtrip [%d]\" i)\n (Float.abs (Nx.item [ i ] d_kpc' -. Nx.item [ i ] d_orig) < 0.01)\n done\n\n(* Cosmo: growth and power spectrum *)\n\nlet test_cosmo_growth_factor_z0 () =\n let g = v (Cosmo.growth_factor ~p:Cosmo.planck18 (Nx.scalar f64 0.0)) in\n is_true\n ~msg:(Printf.sprintf \"D(z=0) = 1.0, got %.6f\" g)\n (Float.abs (g -. 1.0) < 1e-4)\n\nlet test_cosmo_growth_factor_z1 () =\n let g = v (Cosmo.growth_factor ~p:Cosmo.planck18 (Nx.scalar f64 1.0)) in\n is_true\n ~msg:(Printf.sprintf \"D(z=1) ~ 0.61, got %.4f\" g)\n (Float.abs (g -. 0.61) < 0.02)\n\nlet test_cosmo_growth_rate_z0 () =\n let f = v (Cosmo.growth_rate ~p:Cosmo.planck18 (Nx.scalar f64 0.0)) in\n (* f(z=0) ~ Ω_m^0.55 ~ 0.524 for Planck18 *)\n is_true\n ~msg:(Printf.sprintf \"f(z=0) ~ 0.52, got %.4f\" f)\n (Float.abs (f -. 0.52) < 0.02)\n\nlet test_cosmo_growth_monotonic () =\n let p = Cosmo.planck18 in\n let d0 = v (Cosmo.growth_factor ~p (Nx.scalar f64 0.0)) in\n let d1 = v (Cosmo.growth_factor ~p (Nx.scalar f64 0.5)) in\n let d2 = v (Cosmo.growth_factor ~p (Nx.scalar f64 1.0)) in\n is_true ~msg:\"D(0) > D(0.5) > D(1)\" (d0 > d1 && d1 > d2)\n\nlet test_cosmo_linear_power () =\n let p = Cosmo.planck18 in\n let k = Nx.scalar f64 0.1 in\n let pk = v (Cosmo.linear_power ~p k (Nx.scalar f64 0.0)) in\n is_true ~msg:(Printf.sprintf \"P_lin(k=0.1, z=0) > 0, got %.1f\" pk) (pk > 0.0);\n (* P(k, z=1) should be less than P(k, z=0) *)\n let pk1 = v (Cosmo.linear_power ~p k (Nx.scalar f64 1.0)) in\n is_true ~msg:\"P_lin(z=1) < P_lin(z=0)\" (pk1 < pk)\n\nlet test_cosmo_nonlinear_power () =\n let p = Cosmo.planck18 in\n let k = Nx.scalar f64 1.0 in\n let pk_nl = v (Cosmo.nonlinear_power ~p k (Nx.scalar f64 0.0)) in\n let pk_lin = v (Cosmo.linear_power ~p k (Nx.scalar f64 0.0)) in\n is_true ~msg:(Printf.sprintf \"P_nl(k=1) > 0, got %.1f\" pk_nl) (pk_nl > 0.0);\n (* At k=1 h/Mpc, nonlinear should exceed linear *)\n is_true\n ~msg:(Printf.sprintf \"P_nl(k=1) > P_lin(k=1): %.1f > %.1f\" pk_nl pk_lin)\n (pk_nl > pk_lin)\n\nlet test_cosmo_params_accessors () =\n let p = Cosmo.planck18 in\n let ob = v (Cosmo.omega_b p) in\n let ns = v (Cosmo.n_s p) in\n let s8 = v (Cosmo.sigma8 p) in\n is_true ~msg:\"Planck18 omega_b = 0.049\" (Float.abs (ob -. 0.049) < 1e-6);\n is_true ~msg:\"Planck18 n_s = 0.9665\" (Float.abs (ns -. 0.9665) < 1e-6);\n is_true ~msg:\"Planck18 sigma8 = 0.8102\" (Float.abs (s8 -. 0.8102) < 1e-6)\n\n(* Survey tests *)\n\nlet test_survey_smail_normalized () =\n let nz = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.3 () in\n let n = 1000 in\n let zmax = Survey.nz_zmax nz in\n let dz = zmax /. Float.of_int n in\n let sum = ref 0.0 in\n for i = 0 to n do\n let z = Float.of_int i *. dz in\n let nz_val = v (Survey.eval_nz nz (Nx.scalar f64 z)) in\n let w = if i = 0 || i = n then 0.5 else 1.0 in\n sum := !sum +. (w *. nz_val *. dz)\n done;\n is_true\n ~msg:(Printf.sprintf \"smail integrates to 1.0, got %.6f\" !sum)\n (Float.abs (!sum -. 1.0) < 1e-3)\n\nlet test_survey_tabulated () =\n let z = Nx.create f64 [| 5 |] [| 0.0; 0.25; 0.5; 0.75; 1.0 |] in\n let pz = Nx.create f64 [| 5 |] [| 0.0; 1.0; 2.0; 1.0; 0.0 |] in\n let nz = Survey.tabulated ~z ~pz () in\n let mid = v (Survey.eval_nz nz (Nx.scalar f64 0.5)) in\n is_true ~msg:\"tabulated mid > 0\" (mid > 0.0);\n let out = v (Survey.eval_nz nz (Nx.scalar f64 1.5)) in\n is_true ~msg:\"tabulated outside = 0\" (Float.abs out < eps)\n\nlet test_survey_cl_shape () =\n let p = Cosmo.planck18 in\n let nz1 = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.3 () in\n let nz2 = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.7 () in\n let wl1 = Survey.weak_lensing ~n_gal:26.0 nz1 in\n let wl2 = Survey.weak_lensing ~n_gal:26.0 nz2 in\n let ell = Nx.create f64 [| 3 |] [| 100.0; 300.0; 1000.0 |] in\n let cls = Survey.angular_cl ~p ~power:Survey.linear ~ell [ wl1; wl2 ] in\n let shape = Nx.shape (Survey.Cls.to_tensor cls) in\n is_true\n ~msg:(Printf.sprintf \"C_l shape = [3; 3], got [%d; %d]\" shape.(0) shape.(1))\n (shape.(0) = 3 && shape.(1) = 3);\n is_true\n ~msg:(Printf.sprintf \"n_tracers = 2, got %d\" (Survey.Cls.n_tracers cls))\n (Survey.Cls.n_tracers cls = 2)\n\nlet test_survey_cl_positive () =\n let p = Cosmo.planck18 in\n let nz1 = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.5 () in\n let wl = Survey.weak_lensing ~n_gal:26.0 nz1 in\n let ell = Nx.create f64 [| 3 |] [| 100.0; 500.0; 1000.0 |] in\n let cls = Survey.angular_cl ~p ~power:Survey.linear ~ell [ wl ] in\n let cl_auto = Survey.Cls.get cls ~i:0 ~j:0 in\n for l = 0 to 2 do\n let cl_val = Nx.item [ l ] cl_auto in\n is_true ~msg:(Printf.sprintf \"C_l[%d] = %.2e > 0\" l cl_val) (cl_val > 0.0)\n done\n\nlet test_survey_noise_wl () =\n let sigma_e = 0.26 in\n let n_gal = 30.0 in\n let nz1 = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.3 () in\n let wl = Survey.weak_lensing ~sigma_e ~n_gal nz1 in\n let ell = Nx.create f64 [| 3 |] [| 100.0; 500.0; 1000.0 |] in\n let cls = Survey.angular_cl ~ell [ wl ] in\n let nl = Survey.Cls.noise cls in\n let n0 = Nx.item [ 0; 0 ] nl in\n let n1 = Nx.item [ 0; 1 ] nl in\n let n2 = Nx.item [ 0; 2 ] nl in\n is_true ~msg:\"WL noise > 0\" (n0 > 0.0);\n is_true\n ~msg:(Printf.sprintf \"WL noise constant in ℓ: %.2e vs %.2e\" n0 n1)\n (Float.abs (n0 -. n1) < 1e-20);\n is_true ~msg:\"WL noise constant in ℓ (2)\" (Float.abs (n1 -. n2) < 1e-20)\n\n(* Spectrum: mul/div *)\n\nlet test_spectrum_mul () =\n let wave = Unit.Length.of_m (Nx.linspace f64 3e-7 1e-6 10) in\n let values =\n Nx.create f64 [| 10 |]\n [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0; 7.0; 8.0; 9.0; 10.0 |]\n in\n let a =\n Spectrum.create ~wavelength:wave ~values |> Spectrum.as_flux_density\n in\n let trans =\n Nx.create f64 [| 10 |]\n [| 0.5; 0.5; 0.5; 0.5; 0.5; 0.5; 0.5; 0.5; 0.5; 0.5 |]\n in\n let b = Spectrum.create ~wavelength:wave ~values:trans in\n let result = Spectrum.mul a b in\n is_true ~msg:\"mul: 2.0 * 0.5 = 1.0\"\n (Float.abs (Nx.item [ 1 ] (Spectrum.values result) -. 1.0) < eps);\n is_true ~msg:\"mul: 10.0 * 0.5 = 5.0\"\n (Float.abs (Nx.item [ 9 ] (Spectrum.values result) -. 5.0) < eps)\n\nlet test_spectrum_div () =\n let wave = Unit.Length.of_m (Nx.linspace f64 3e-7 1e-6 10) in\n let values =\n Nx.create f64 [| 10 |]\n [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0; 7.0; 8.0; 9.0; 10.0 |]\n in\n let a =\n Spectrum.create ~wavelength:wave ~values |> Spectrum.as_flux_density\n in\n let flat =\n Nx.create f64 [| 10 |]\n [| 2.0; 2.0; 2.0; 2.0; 2.0; 2.0; 2.0; 2.0; 2.0; 2.0 |]\n in\n let b = Spectrum.create ~wavelength:wave ~values:flat in\n let result = Spectrum.div a b in\n is_true ~msg:\"div: 4.0 / 2.0 = 2.0\"\n (Float.abs (Nx.item [ 3 ] (Spectrum.values result) -. 2.0) < eps);\n is_true ~msg:\"div: 10.0 / 2.0 = 5.0\"\n (Float.abs (Nx.item [ 9 ] (Spectrum.values result) -. 5.0) < eps)\n\nlet test_spectrum_mul_div_roundtrip () =\n let wave = Unit.Length.of_m (Nx.linspace f64 3e-7 1e-6 50) in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 6000.0) in\n let spec =\n Spectrum.blackbody ~temperature:temp ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let trans_vals =\n Nx.create f64 [| 50 |]\n (Array.init 50 (fun i ->\n 0.5 +. (0.3 *. Float.sin (Float.of_int i *. 0.2))))\n in\n let trans = Spectrum.create ~wavelength:wave ~values:trans_vals in\n let mulled = Spectrum.mul spec trans in\n let recovered = Spectrum.div mulled trans in\n let orig_val = Nx.item [ 25 ] (Spectrum.values spec) in\n let rec_val = Nx.item [ 25 ] (Spectrum.values recovered) in\n is_true ~msg:\"mul then div roundtrip\"\n (Float.abs (rec_val -. orig_val) /. orig_val < 1e-10)\n\n(* Spectrum: line profiles *)\n\nlet test_spectrum_gaussian_peak () =\n let wave = Unit.Length.of_m (Nx.linspace f64 6.4e-7 6.7e-7 1000) in\n let center = Unit.Length.nm 656.3 in\n let stddev = Unit.Length.nm 1.0 in\n let amplitude = Nx.scalar f64 1.0 in\n let g = Spectrum.gaussian ~amplitude ~center ~stddev ~wavelength:wave in\n let vals = Spectrum.values g in\n let peak_idx = ref 0 in\n let peak_val = ref (Nx.item [ 0 ] vals) in\n for i = 1 to 999 do\n let vi = Nx.item [ i ] vals in\n if vi > !peak_val then begin\n peak_val := vi;\n peak_idx := i\n end\n done;\n let wave_m = Unit.Length.in_m (Spectrum.wavelength g) in\n let peak_lam_nm = Nx.item [ !peak_idx ] wave_m *. 1e9 in\n is_true\n ~msg:(Printf.sprintf \"Gaussian peak near 656.3 nm, got %.1f\" peak_lam_nm)\n (Float.abs (peak_lam_nm -. 656.3) < 0.5);\n is_true\n ~msg:(Printf.sprintf \"Gaussian peak amplitude ~ 1.0, got %.4f\" !peak_val)\n (Float.abs (!peak_val -. 1.0) < 0.01)\n\nlet test_spectrum_lorentzian_peak () =\n let wave = Unit.Length.of_m (Nx.linspace f64 4.8e-7 5.2e-7 1000) in\n let center = Unit.Length.nm 500.0 in\n let fwhm = Unit.Length.nm 2.0 in\n let amplitude = Nx.scalar f64 3.0 in\n let l = Spectrum.lorentzian ~amplitude ~center ~fwhm ~wavelength:wave in\n let vals = Spectrum.values l in\n let peak_idx = ref 0 in\n let peak_val = ref (Nx.item [ 0 ] vals) in\n for i = 1 to 999 do\n let vi = Nx.item [ i ] vals in\n if vi > !peak_val then begin\n peak_val := vi;\n peak_idx := i\n end\n done;\n let wave_m = Unit.Length.in_m (Spectrum.wavelength l) in\n let peak_lam_nm = Nx.item [ !peak_idx ] wave_m *. 1e9 in\n is_true\n ~msg:(Printf.sprintf \"Lorentzian peak near 500 nm, got %.1f\" peak_lam_nm)\n (Float.abs (peak_lam_nm -. 500.0) < 0.5);\n is_true\n ~msg:(Printf.sprintf \"Lorentzian peak ~ 3.0, got %.4f\" !peak_val)\n (Float.abs (!peak_val -. 3.0) < 0.05)\n\nlet test_spectrum_voigt_limits () =\n let wave = Unit.Length.of_m (Nx.linspace f64 4.8e-7 5.2e-7 1000) in\n let center = Unit.Length.nm 500.0 in\n let amplitude = Nx.scalar f64 1.0 in\n (* Gaussian limit: sigma >> gamma *)\n let sigma_big = Unit.Length.nm 2.0 in\n let gamma_tiny = Unit.Length.nm 0.001 in\n let voigt_g =\n Spectrum.voigt ~amplitude ~center ~sigma:sigma_big ~gamma:gamma_tiny\n ~wavelength:wave\n in\n let gauss =\n Spectrum.gaussian ~amplitude ~center ~stddev:sigma_big ~wavelength:wave\n in\n let vg_peak = ref 0.0 in\n let g_peak = ref 0.0 in\n for i = 0 to 999 do\n let vv = Nx.item [ i ] (Spectrum.values voigt_g) in\n let gv = Nx.item [ i ] (Spectrum.values gauss) in\n if vv > !vg_peak then vg_peak := vv;\n if gv > !g_peak then g_peak := gv\n done;\n is_true\n ~msg:\n (Printf.sprintf \"Voigt(sigma>>gamma) peak ~ Gaussian peak: %.4f vs %.4f\"\n !vg_peak !g_peak)\n (Float.abs (!vg_peak -. !g_peak) /. !g_peak < 0.05)\n\nlet test_spectrum_line_composability () =\n let wave = Unit.Length.of_m (Nx.linspace f64 6e-7 7e-7 500) in\n let continuum =\n Spectrum.power_law ~amplitude:(Nx.scalar f64 1e-15)\n ~index:(Nx.scalar f64 (-2.0)) ~pivot:(Unit.Length.nm 650.0)\n ~wavelength:wave\n in\n let ha =\n Spectrum.gaussian ~amplitude:(Nx.scalar f64 1e-15)\n ~center:(Unit.Length.nm 656.3) ~stddev:(Unit.Length.nm 0.5)\n ~wavelength:wave\n in\n let composite = Spectrum.add continuum ha in\n let cont_val = Nx.item [ 0 ] (Spectrum.values continuum) in\n let comp_val = Nx.item [ 0 ] (Spectrum.values composite) in\n is_true ~msg:\"Composite spectrum at wing ~ continuum\"\n (Float.abs (comp_val -. cont_val) /. cont_val < 0.01)\n\n(* Altaz: airmass *)\n\nlet test_altaz_airmass_zenith () =\n let hz =\n Altaz.of_coord\n ~obstime:(Time.of_iso \"2024-06-21T12:00:00\")\n ~observer:\n (Altaz.make_observer ~lat:(Unit.Angle.deg 45.0)\n ~lon:(Unit.Angle.deg 0.0) ())\n (Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 0.0 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 89.0 |])))\n in\n let x = Altaz.airmass hz in\n let x0 = Nx.item [ 0 ] x in\n is_true ~msg:(Printf.sprintf \"Airmass >= 1.0, got %.4f\" x0) (x0 >= 1.0)\n\nlet test_altaz_airmass_low_alt () =\n let obs =\n Altaz.make_observer ~lat:(Unit.Angle.deg 30.0) ~lon:(Unit.Angle.deg 0.0) ()\n in\n let t = Time.of_iso \"2024-06-21T22:00:00\" in\n let star_a =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 0.0 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 80.0 |]))\n in\n let star_b =\n Coord.of_radec\n ~ra:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 180.0 |]))\n ~dec:(Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 10.0 |]))\n in\n let hz_a = Altaz.of_coord ~obstime:t ~observer:obs star_a in\n let hz_b = Altaz.of_coord ~obstime:t ~observer:obs star_b in\n let x_a = Nx.item [ 0 ] (Altaz.airmass hz_a) in\n let x_b = Nx.item [ 0 ] (Altaz.airmass hz_b) in\n is_true\n ~msg:(Printf.sprintf \"Both airmasses >= 1: %.2f, %.2f\" x_a x_b)\n (x_a >= 1.0 && x_b >= 1.0);\n is_true\n ~msg:(Printf.sprintf \"Different airmasses: %.2f vs %.2f\" x_a x_b)\n (Float.abs (x_a -. x_b) > 0.01)\n\n(* Cosmo: BAO distances *)\n\nlet test_cosmo_dh () =\n let p = Cosmo.planck18 in\n let z = Nx.scalar f64 0.0 in\n let dh0 = v (Unit.Length.in_mpc (Cosmo.dh ~p z)) in\n let h0 = Nx.item [] (Cosmo.h0 p) in\n let expected = 299792.458 /. h0 in\n is_true\n ~msg:(Printf.sprintf \"D_H(0) = c/H0 ~ %.1f Mpc, got %.1f\" expected dh0)\n (Float.abs (dh0 -. expected) /. expected < 1e-4)\n\nlet test_cosmo_dm_flat () =\n let p = Cosmo.planck18 in\n let z = Nx.scalar f64 0.5 in\n let dm_val = v (Unit.Length.in_mpc (Cosmo.dm ~p z)) in\n let dc_val = v (Unit.Length.in_mpc (Cosmo.comoving_distance ~p z)) in\n is_true\n ~msg:(Printf.sprintf \"D_M = D_C for flat: %.1f vs %.1f\" dm_val dc_val)\n (Float.abs (dm_val -. dc_val) /. dc_val < 1e-4)\n\nlet test_cosmo_dv () =\n let p = Cosmo.planck18 in\n let z = Nx.scalar f64 0.5 in\n let dv_val = v (Unit.Length.in_mpc (Cosmo.dv ~p z)) in\n is_true ~msg:(Printf.sprintf \"D_V(0.5) > 0, got %.1f\" dv_val) (dv_val > 0.0);\n let dh_val = v (Unit.Length.in_mpc (Cosmo.dh ~p z)) in\n let dm_val = v (Unit.Length.in_mpc (Cosmo.dm ~p z)) in\n let z_f = 0.5 in\n let expected = (z_f *. dh_val *. dm_val *. dm_val) ** (1.0 /. 3.0) in\n is_true\n ~msg:\n (Printf.sprintf \"D_V = (z D_H D_M^2)^{1/3}: %.1f vs %.1f\" dv_val expected)\n (Float.abs (dv_val -. expected) /. expected < 1e-3)\n\nlet test_cosmo_sound_horizon () =\n let p = Cosmo.planck18 in\n let rs = v (Unit.Length.in_mpc (Cosmo.sound_horizon ~p ())) in\n is_true\n ~msg:(Printf.sprintf \"r_s(Planck18) ~ 147 Mpc, got %.1f\" rs)\n (Float.abs (rs -. 147.0) < 5.0)\n\n(* Filters *)\n\nlet test_filters_sdss_pivot () =\n let bp = Filters.sdss_r in\n let lam_p = v (Unit.Length.in_nm (Photometry.pivot_wavelength bp)) in\n is_true\n ~msg:(Printf.sprintf \"SDSS r pivot ~ 620 nm, got %.0f\" lam_p)\n (Float.abs (lam_p -. 620.0) < 30.0)\n\nlet test_filters_johnson_v_pivot () =\n let bp = Filters.johnson_v in\n let lam_p = v (Unit.Length.in_nm (Photometry.pivot_wavelength bp)) in\n is_true\n ~msg:(Printf.sprintf \"Johnson V pivot ~ 551 nm, got %.0f\" lam_p)\n (Float.abs (lam_p -. 551.0) < 20.0)\n\nlet test_filters_twomass_j_pivot () =\n let bp = Filters.twomass_j in\n let lam_p = v (Unit.Length.in_nm (Photometry.pivot_wavelength bp)) in\n is_true\n ~msg:(Printf.sprintf \"2MASS J pivot ~ 1235 nm, got %.0f\" lam_p)\n (Float.abs (lam_p -. 1235.0) < 30.0)\n\nlet test_filters_gaia_ordering () =\n let bp_p =\n v (Unit.Length.in_nm (Photometry.pivot_wavelength Filters.gaia_bp))\n in\n let g_p =\n v (Unit.Length.in_nm (Photometry.pivot_wavelength Filters.gaia_g))\n in\n let rp_p =\n v (Unit.Length.in_nm (Photometry.pivot_wavelength Filters.gaia_rp))\n in\n is_true\n ~msg:(Printf.sprintf \"Gaia: BP < G < RP: %.0f < %.0f < %.0f\" bp_p g_p rp_p)\n (bp_p < g_p && g_p < rp_p)\n\nlet test_filters_photometry () =\n let bp = Filters.sdss_g in\n let wave = Photometry.wavelength bp in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 5800.0) in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let mag = Nx.item [] (Photometry.ab_mag bp sed) in\n is_true\n ~msg:(Printf.sprintf \"BB(5800K) through SDSS g is finite, got %.2f\" mag)\n (Float.is_finite mag)\n\n(* Photometry: auto-resample *)\n\nlet test_photometry_auto_resample () =\n let bp = Filters.sdss_g in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 5800.0) in\n let wave_fine = Unit.Length.of_m (Nx.linspace f64 3e-7 1.1e-6 1000) in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:wave_fine\n |> Spectrum.as_flux_density\n in\n let mag = Nx.item [] (Photometry.ab_mag bp sed) in\n is_true\n ~msg:\n (Printf.sprintf \"Auto-resample: BB(5800K) through SDSS g finite, got %.2f\"\n mag)\n (Float.is_finite mag);\n let manual = Spectrum.resample ~wavelength:(Photometry.wavelength bp) sed in\n let mag_manual = Nx.item [] (Photometry.ab_mag bp manual) in\n is_true\n ~msg:\n (Printf.sprintf \"Auto-resample matches manual: %.4f vs %.4f\" mag\n mag_manual)\n (Float.abs (mag -. mag_manual) < 1e-10)\n\n(* Photometry: ST magnitude *)\n\nlet test_photometry_st_mag () =\n let bp =\n Photometry.tophat ~lo:(Unit.Length.nm 400.0) ~hi:(Unit.Length.nm 700.0)\n ~n:100\n in\n let wave = Photometry.wavelength bp in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 5800.0) in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let st = Nx.item [] (Photometry.st_mag bp sed) in\n let ab = Nx.item [] (Photometry.ab_mag bp sed) in\n is_true ~msg:(Printf.sprintf \"ST mag is finite: %.2f\" st) (Float.is_finite st);\n is_true\n ~msg:(Printf.sprintf \"ST and AB differ: ST=%.2f AB=%.2f\" st ab)\n (Float.abs (st -. ab) > 0.01)\n\n(* Photometry: Vega magnitude *)\n\nlet test_photometry_vega_mag () =\n let bp = Filters.johnson_v in\n let wave = Photometry.wavelength bp in\n let temp = Unit.Temperature.of_kelvin (Nx.scalar f64 9600.0) in\n let sed =\n Spectrum.blackbody ~temperature:temp ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let vm = Nx.item [] (Photometry.vega_mag bp sed) in\n is_true\n ~msg:(Printf.sprintf \"Vega mag of hot BB through V is finite: %.2f\" vm)\n (Float.is_finite vm);\n let ab = Nx.item [] (Photometry.ab_mag bp sed) in\n is_true\n ~msg:(Printf.sprintf \"Vega and AB differ: V=%.2f AB=%.2f\" vm ab)\n (Float.abs (vm -. ab) > 0.001)\n\n(* Photometry: effective wavelength *)\n\nlet test_photometry_effective_wavelength () =\n let bp =\n Photometry.tophat ~lo:(Unit.Length.nm 400.0) ~hi:(Unit.Length.nm 700.0)\n ~n:100\n in\n let wave = Photometry.wavelength bp in\n let flat_vals = Nx.ones f64 [| 100 |] in\n let flat =\n Spectrum.create ~wavelength:wave ~values:flat_vals\n |> Spectrum.as_flux_density\n in\n let lam_eff =\n v (Unit.Length.in_nm (Photometry.effective_wavelength bp flat))\n in\n let lam_pivot = v (Unit.Length.in_nm (Photometry.pivot_wavelength bp)) in\n is_true\n ~msg:\n (Printf.sprintf \"Flat spectrum: eff_wavelength in range: %.1f nm\" lam_eff)\n (lam_eff > 500.0 && lam_eff < 600.0);\n is_true\n ~msg:\n (Printf.sprintf \"eff_wavelength >= pivot for flat/tophat: %.1f vs %.1f\"\n lam_eff lam_pivot)\n (lam_eff >= lam_pivot)\n\n(* Altaz: atmospheric refraction *)\n\nlet test_altaz_refraction () =\n let obs =\n Altaz.make_observer ~lat:(Unit.Angle.deg 45.0) ~lon:(Unit.Angle.deg 0.0) ()\n in\n let t = Time.of_iso \"2024-06-15T12:00:00\" in\n let ra = Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 37.95 |]) in\n let dec = Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 89.264 |]) in\n let c = Coord.of_radec ~ra ~dec in\n let hz_no = Altaz.of_coord ~refraction:false ~obstime:t ~observer:obs c in\n let hz_yes = Altaz.of_coord ~refraction:true ~obstime:t ~observer:obs c in\n let alt_no = Nx.item [ 0 ] (Unit.Angle.in_deg (Altaz.alt hz_no)) in\n let alt_yes = Nx.item [ 0 ] (Unit.Angle.in_deg (Altaz.alt hz_yes)) in\n (* Refraction makes objects appear higher *)\n is_true\n ~msg:\n (Printf.sprintf \"Refraction raises altitude: %.4f > %.4f\" alt_yes alt_no)\n (alt_yes > alt_no);\n (* At ~45° alt, refraction is ~1 arcmin = 0.017° *)\n let diff = alt_yes -. alt_no in\n is_true\n ~msg:(Printf.sprintf \"Refraction at ~45° is small (< 0.1°): %.4f\" diff)\n (diff > 0.0 && diff < 0.1)\n\nlet test_altaz_refraction_standalone () =\n let obs =\n Altaz.make_observer ~lat:(Unit.Angle.deg 45.0) ~lon:(Unit.Angle.deg 0.0) ()\n in\n let t = Time.of_iso \"2024-06-15T12:00:00\" in\n let ra = Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 37.95 |]) in\n let dec = Unit.Angle.of_deg (Nx.create f64 [| 1 |] [| 89.264 |]) in\n let c = Coord.of_radec ~ra ~dec in\n let hz = Altaz.of_coord ~obstime:t ~observer:obs c in\n let r = Altaz.refraction hz in\n let r_arcmin = Nx.item [ 0 ] (Unit.Angle.in_deg r) *. 60.0 in\n is_true\n ~msg:(Printf.sprintf \"Refraction > 0 arcmin: %.2f\" r_arcmin)\n (r_arcmin > 0.0);\n is_true\n ~msg:(Printf.sprintf \"Refraction < 2 arcmin at high alt: %.2f\" r_arcmin)\n (r_arcmin < 2.0)\n\n(* Survey: shear multiplicative bias *)\n\nlet test_survey_m_bias () =\n let nz = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.5 () in\n let ell = Nx.logspace f64 1.0 3.0 20 in\n let wl_no_bias = Survey.weak_lensing ~n_gal:26.0 nz in\n let wl_with_bias = Survey.weak_lensing ~m_bias:0.02 ~n_gal:26.0 nz in\n let cls_no = Survey.angular_cl ~ell [ wl_no_bias ] in\n let cls_yes = Survey.angular_cl ~ell [ wl_with_bias ] in\n let cl_no = Survey.Cls.get cls_no ~i:0 ~j:0 in\n let cl_yes = Survey.Cls.get cls_yes ~i:0 ~j:0 in\n (* Auto-spectrum scales as (1+m)^2 = 1.0404 *)\n let ratio = Nx.item [ 10 ] (Nx.div cl_yes cl_no) in\n let expected = 1.02 *. 1.02 in\n is_true\n ~msg:\n (Printf.sprintf\n \"m_bias=0.02 scales auto-Cl by (1+m)^2: ratio=%.4f vs %.4f\" ratio\n expected)\n (Float.abs (ratio -. expected) < 1e-4);\n (* m_bias=0.0 gives same result as no bias *)\n let wl_zero_bias = Survey.weak_lensing ~m_bias:0.0 ~n_gal:26.0 nz in\n let cls_zero = Survey.angular_cl ~ell [ wl_zero_bias ] in\n let cl_zero = Survey.Cls.get cls_zero ~i:0 ~j:0 in\n let diff = Nx.item [] (Nx.max (Nx.abs (Nx.sub cl_zero cl_no))) in\n is_true\n ~msg:(Printf.sprintf \"m_bias=0.0 matches no bias: max_diff=%.2e\" diff)\n (diff < 1e-30)\n\n(* Spectrum: differentiable resample *)\n\nlet test_spectrum_resample_values () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-5 100) in\n let sed =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 5800.0))\n ~wavelength:wave\n in\n let new_wave = Unit.Length.of_m (Nx.linspace f64 2e-7 9e-6 50) in\n let resampled = Spectrum.resample ~wavelength:new_wave sed in\n let vals = Spectrum.values resampled in\n let n = Nx.numel vals in\n is_true ~msg:(Printf.sprintf \"Resampled has %d points\" n) (n = 50);\n let v0 = Nx.item [ 0 ] vals in\n is_true ~msg:(Printf.sprintf \"Resampled values positive: %.2e\" v0) (v0 > 0.0);\n let vmax = Nx.item [] (Nx.max vals) in\n is_true\n ~msg:(Printf.sprintf \"Resampled max is finite: %.2e\" vmax)\n (Float.is_finite vmax)\n\n(* Survey: baryonic feedback *)\n\nlet test_survey_baryonic_feedback () =\n let nz = Survey.smail ~a:2.0 ~b:1.5 ~z0:0.5 () in\n let ell = Nx.logspace f64 1.0 3.0 20 in\n let wl = Survey.weak_lensing ~n_gal:26.0 nz in\n let cls_dm = Survey.angular_cl ~power:Survey.nonlinear ~ell [ wl ] in\n let power_bary = Survey.baryonic_feedback ~a_bary:0.2 Survey.nonlinear in\n let cls_bary = Survey.angular_cl ~power:power_bary ~ell [ wl ] in\n let cl_dm = Survey.Cls.get cls_dm ~i:0 ~j:0 in\n let cl_bary = Survey.Cls.get cls_bary ~i:0 ~j:0 in\n (* Baryonic feedback suppresses small-scale (high-ell) power *)\n let ratio_high = Nx.item [ 19 ] (Nx.div cl_bary cl_dm) in\n is_true\n ~msg:\n (Printf.sprintf \"Baryonic suppression at high ell: ratio=%.4f < 1\"\n ratio_high)\n (ratio_high < 1.0);\n (* a_bary=0 gives same result as no feedback *)\n let power_zero = Survey.baryonic_feedback ~a_bary:0.0 Survey.nonlinear in\n let cls_zero = Survey.angular_cl ~power:power_zero ~ell [ wl ] in\n let cl_zero = Survey.Cls.get cls_zero ~i:0 ~j:0 in\n let diff = Nx.item [] (Nx.max (Nx.abs (Nx.sub cl_zero cl_dm))) in\n is_true\n ~msg:(Printf.sprintf \"a_bary=0 matches DM-only: max_diff=%.2e\" diff)\n (diff < 1e-30)\n\n(* Batched spectra *)\n\nlet test_batch_create () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-6 100) in\n let v1 = Nx.ones f64 [| 100 |] in\n let v2 = Nx.full f64 [| 100 |] 2.0 in\n let values = Nx.stack [ v1; v2 ] in\n let s =\n Spectrum.create ~wavelength:wave ~values |> Spectrum.as_flux_density\n in\n let sh = Nx.shape (Spectrum.values s) in\n is_true ~msg:\"batch values shape [2; 100]\"\n (Array.length sh = 2 && sh.(0) = 2 && sh.(1) = 100)\n\nlet test_batch_resample () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-6 100) in\n let wave2 = Unit.Length.of_m (Nx.linspace f64 2e-7 8e-7 50) in\n let bb1 =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 5000.0))\n ~wavelength:wave\n |> Spectrum.as_sampled\n in\n let bb2 =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 8000.0))\n ~wavelength:wave\n |> Spectrum.as_sampled\n in\n let values = Nx.stack [ Spectrum.values bb1; Spectrum.values bb2 ] in\n let batch = Spectrum.create ~wavelength:wave ~values |> Spectrum.as_sampled in\n let resampled = Spectrum.resample ~wavelength:wave2 batch in\n let r_shape = Nx.shape (Spectrum.values resampled) in\n is_true ~msg:\"batch resample shape [2; 50]\"\n (Array.length r_shape = 2 && r_shape.(0) = 2 && r_shape.(1) = 50);\n let r1 = Spectrum.resample ~wavelength:wave2 bb1 in\n let r2 = Spectrum.resample ~wavelength:wave2 bb2 in\n let expected = Nx.stack [ Spectrum.values r1; Spectrum.values r2 ] in\n let diff =\n Nx.item [] (Nx.max (Nx.abs (Nx.sub (Spectrum.values resampled) expected)))\n in\n is_true\n ~msg:\n (Printf.sprintf \"batch resample matches individual: max_diff=%.2e\" diff)\n (diff < 1e-20)\n\nlet test_batch_ab_mag () =\n let wave = Unit.Length.of_m (Nx.linspace f64 3e-7 8e-7 200) in\n let bp =\n Photometry.tophat ~lo:(Unit.Length.nm 400.0) ~hi:(Unit.Length.nm 600.0)\n ~n:100\n in\n let bb1 =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 5000.0))\n ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let bb2 =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 8000.0))\n ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let values = Nx.stack [ Spectrum.values bb1; Spectrum.values bb2 ] in\n let batch =\n Spectrum.create ~wavelength:wave ~values |> Spectrum.as_flux_density\n in\n let mags_batch = Photometry.ab_mag bp batch in\n let mag1 = Photometry.ab_mag bp bb1 in\n let mag2 = Photometry.ab_mag bp bb2 in\n let expected = Nx.stack [ mag1; mag2 ] in\n let diff = Nx.item [] (Nx.max (Nx.abs (Nx.sub mags_batch expected))) in\n is_true\n ~msg:(Printf.sprintf \"batch ab_mag matches individual: max_diff=%.2e\" diff)\n (diff < 1e-10)\n\nlet test_batch_extinction () =\n let wave = Unit.Length.of_m (Nx.linspace f64 3e-7 8e-7 200) in\n let rv = Nx.scalar f64 3.1 in\n let av = Nx.scalar f64 0.5 in\n let bb1 =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 5000.0))\n ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let bb2 =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 8000.0))\n ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let values = Nx.stack [ Spectrum.values bb1; Spectrum.values bb2 ] in\n let batch =\n Spectrum.create ~wavelength:wave ~values |> Spectrum.as_flux_density\n in\n let reddened = Extinction.apply (Extinction.ccm89 ~rv) ~av batch in\n let r1 = Extinction.apply (Extinction.ccm89 ~rv) ~av bb1 in\n let r2 = Extinction.apply (Extinction.ccm89 ~rv) ~av bb2 in\n let expected = Nx.stack [ Spectrum.values r1; Spectrum.values r2 ] in\n let diff =\n Nx.item [] (Nx.max (Nx.abs (Nx.sub (Spectrum.values reddened) expected)))\n in\n is_true\n ~msg:\n (Printf.sprintf \"batch extinction matches individual: max_diff=%.2e\" diff)\n (diff < 1e-25)\n\nlet test_batch_scale () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-6 100) in\n let v1 = Nx.ones f64 [| 100 |] in\n let v2 = Nx.full f64 [| 100 |] 2.0 in\n let values = Nx.stack [ v1; v2 ] in\n let batch = Spectrum.create ~wavelength:wave ~values in\n let scaled = Spectrum.scale (Nx.scalar f64 3.0) batch in\n let sv = Spectrum.values scaled in\n let expected =\n Nx.stack [ Nx.full f64 [| 100 |] 3.0; Nx.full f64 [| 100 |] 6.0 ]\n in\n let diff = Nx.item [] (Nx.max (Nx.abs (Nx.sub sv expected))) in\n is_true ~msg:\"batch scalar scale\" (diff < 1e-15)\n\nlet test_batch_redshift_scalar () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-6 100) in\n let bb1 =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 5000.0))\n ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let bb2 =\n Spectrum.blackbody\n ~temperature:(Unit.Temperature.of_kelvin (Nx.scalar f64 8000.0))\n ~wavelength:wave\n |> Spectrum.as_flux_density\n in\n let values = Nx.stack [ Spectrum.values bb1; Spectrum.values bb2 ] in\n let batch =\n Spectrum.create ~wavelength:wave ~values |> Spectrum.as_flux_density\n in\n let z = Nx.scalar f64 0.5 in\n let shifted = Spectrum.redshift ~z batch in\n let s1 = Spectrum.redshift ~z bb1 in\n let s2 = Spectrum.redshift ~z bb2 in\n let expected = Nx.stack [ Spectrum.values s1; Spectrum.values s2 ] in\n let diff =\n Nx.item [] (Nx.max (Nx.abs (Nx.sub (Spectrum.values shifted) expected)))\n in\n is_true\n ~msg:\n (Printf.sprintf \"batch redshift matches individual: max_diff=%.2e\" diff)\n (diff < 1e-20)\n\nlet test_batch_create_mismatch () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-6 100) in\n let values = Nx.ones f64 [| 3; 50 |] in\n let raised =\n try\n ignore (Spectrum.create ~wavelength:wave ~values);\n false\n with Invalid_argument _ -> true\n in\n is_true ~msg:\"mismatched last dim raises\" raised\n\nlet test_batch_roundtrip () =\n let wave = Unit.Length.of_m (Nx.linspace f64 1e-7 1e-6 100) in\n let v1 = Nx.ones f64 [| 100 |] in\n let v2 = Nx.full f64 [| 100 |] 2.0 in\n let v3 = Nx.full f64 [| 100 |] 3.0 in\n let values = Nx.stack [ v1; v2; v3 ] in\n let batch = Spectrum.create ~wavelength:wave ~values in\n let extracted = Nx.get [ 1 ] (Spectrum.values batch) in\n let diff = Nx.item [] (Nx.max (Nx.abs (Nx.sub extracted v2))) in\n is_true ~msg:\"extract second spectrum from batch\" (diff < 1e-15)\n\nlet () =\n run \"Umbra\"\n [\n group \"Unit\"\n [\n test \"10 kpc converts to 0.01 Mpc\" test_length_conversion;\n test \"10 kpc + 500 pc = 10.5 kpc\" test_length_arithmetic;\n test \"1 solar mass is ~1.988e30 kg\" test_mass_conversion;\n test \"100 km / 10 s = 10 km/s\" test_velocity_cross_dim;\n test \"sin(90) = 1 and cos(90) = 0\" test_angle_trig;\n test \"wavelength to frequency roundtrips\" test_wavelength_frequency;\n test \"phantom types prevent adding length and mass\"\n test_phantom_type_safety;\n test \"2 eV survives energy-wavelength-frequency roundtrip\"\n test_energy_wavelength_frequency;\n ];\n group \"Const\" [ test \"speed of light is ~299792 km/s\" test_const_c ];\n group \"Time\"\n [\n test \"J2000.0 JD and MJD values are correct\" test_time_jd_mjd;\n test \"ISO 8601 parse and format roundtrip\" test_time_iso;\n test \"UTC to TAI offset is 32s at J2000\" test_time_utc_tai_tt;\n test \"TDB-TT difference is less than 2 ms\" test_time_tdb;\n test \"Unix epoch maps to JD 2440587.5\" test_time_unix;\n test \"diff and add with 1-day offset\" test_time_diff_add;\n ];\n group \"Coord\"\n [\n test \"ICRS to Galactic and back preserves RA/Dec\" test_coord_roundtrip;\n test \"ICRS to Ecliptic and back preserves RA/Dec\"\n test_coord_ecliptic_roundtrip;\n test \"ICRS to Supergalactic and back preserves RA/Dec\"\n test_coord_supergalactic_roundtrip;\n test \"north pole to south pole separation is 180 deg\"\n test_separation_poles;\n test \"nearest self-match returns identity indices\"\n test_match_nearest_self;\n ];\n group \"Cosmo\"\n [\n test \"H(0) equals H0 = 70 km/s/Mpc\" test_cosmo_hubble;\n test \"E(z=0) = 1 for any cosmology\" test_cosmo_e_of;\n test \"comoving(0.1) ~ 421 Mpc and luminosity(0.1) ~ 463 Mpc\"\n test_cosmo_distances;\n test \"lookback time at z=1 is ~7.7 Gyr\" test_cosmo_lookback;\n test \"1 kpc at z=0.022 subtends ~2.3 arcsec\" test_cosmo_angular_scale;\n test \"Planck18 comoving(0.5) ~ 1960 Mpc\" test_cosmo_planck18;\n test \"flat_lcdm(70, 0.3) matches default cosmology\"\n test_cosmo_flat_lcdm_same_as_default;\n test \"non-flat LCDM differs from flat\" test_cosmo_nonflat_lcdm;\n test \"wCDM with w0=-1 reduces to LCDM\" test_cosmo_wcdm;\n test \"w0waCDM with w0=-1 wa=0 reduces to LCDM\" test_cosmo_w0wacdm;\n test \"z_at_value roundtrips luminosity distance\" test_cosmo_z_at_value;\n test \"growth factor D(z=0) = 1\" test_cosmo_growth_factor_z0;\n test \"growth factor D(z=1) ~ 0.61\" test_cosmo_growth_factor_z1;\n test \"growth rate f(z=0) ~ 0.52\" test_cosmo_growth_rate_z0;\n test \"growth factor decreases with redshift\"\n test_cosmo_growth_monotonic;\n test \"linear power spectrum is positive and decreases with z\"\n test_cosmo_linear_power;\n test \"nonlinear power exceeds linear at k=1 h/Mpc\"\n test_cosmo_nonlinear_power;\n test \"Planck18 omega_b, n_s, and sigma8 accessors\"\n test_cosmo_params_accessors;\n test \"D_H(0) = c/H0\" test_cosmo_dh;\n test \"D_M equals D_C for flat geometry\" test_cosmo_dm_flat;\n test \"D_V = (z * D_H * D_M^2)^(1/3)\" test_cosmo_dv;\n test \"sound horizon r_s ~ 147 Mpc for Planck18\"\n test_cosmo_sound_horizon;\n test \"age of universe ~ 13.8 Gyr for Planck18\" test_cosmo_age_planck18;\n test \"age(z=0) - age(z=1) = lookback(z=1)\" test_cosmo_age_at_z1;\n test \"comoving distance to CMB ~ 14000 Mpc\" test_cosmo_comoving_cmb;\n test \"comoving distances increase at z = 2, 5, 10\"\n test_cosmo_comoving_high_z;\n test \"lookback time at z=5 ~ 12.5 Gyr\" test_cosmo_lookback_high_z;\n ];\n group \"Altaz\"\n [\n test \"ICRS to AltAz and back preserves RA/Dec\" test_altaz_zenith;\n test \"Polaris altitude ~ observer latitude from lat=45\"\n test_altaz_north_pole;\n test \"airmass is >= 1.0 near zenith\" test_altaz_airmass_zenith;\n test \"airmass differs for high vs low altitude stars\"\n test_altaz_airmass_low_alt;\n test \"refraction raises apparent altitude\" test_altaz_refraction;\n test \"standalone refraction is between 0 and 2 arcmin at high alt\"\n test_altaz_refraction_standalone;\n ];\n group \"Galactocentric\"\n [\n test \"l=0 b=0 at galcen_distance maps to origin\"\n test_galactocentric_gc_position;\n test \"Galactocentric to ICRS roundtrips RA/Dec/distance\"\n test_galactocentric_roundtrip;\n ];\n group \"Spectrum\"\n [\n test \"scale by 3 multiplies all values\" test_spectrum_scale;\n test \"multiply spectrum by transmission\" test_spectrum_mul;\n test \"divide spectrum by flat transmission\" test_spectrum_div;\n test \"mul then div roundtrips to original\"\n test_spectrum_mul_div_roundtrip;\n test \"blackbody peak obeys Wien's displacement law\"\n test_spectrum_blackbody_wien;\n test \"redshift z=1 doubles wavelength and halves flux\"\n test_spectrum_redshift;\n test \"resample preserves positivity and finiteness\"\n test_spectrum_resample_values;\n test \"Gaussian line peaks at 656.3 nm with unit amplitude\"\n test_spectrum_gaussian_peak;\n test \"Lorentzian line peaks at 500 nm with amplitude 3\"\n test_spectrum_lorentzian_peak;\n test \"Voigt with sigma >> gamma matches Gaussian\"\n test_spectrum_voigt_limits;\n test \"power-law continuum plus Gaussian line composes cleanly\"\n test_spectrum_line_composability;\n ];\n group \"Extinction\"\n [\n test \"CCM89 A_V/A_V ~ 1.0 at V-band 550 nm\"\n test_extinction_ccm89_v_band;\n test \"CCM89 extinction increases toward blue\"\n test_extinction_ccm89_monotonic;\n test \"apply then unredden recovers original spectrum\"\n test_extinction_apply_unredden;\n ];\n group \"Photometry\"\n [\n test \"flat f_nu = 3631 Jy gives m_AB ~ 0\" test_photometry_ab_mag_flat;\n test \"same-band color is zero\" test_photometry_color_same_band;\n test \"hot star is brighter in blue than red\"\n test_photometry_blue_star_color;\n test \"auto-resample matches manual resample\"\n test_photometry_auto_resample;\n test \"ST and AB magnitudes differ for a blackbody\"\n test_photometry_st_mag;\n test \"Vega and AB magnitudes differ through Johnson V\"\n test_photometry_vega_mag;\n test \"effective wavelength is in range for flat tophat spectrum\"\n test_photometry_effective_wavelength;\n ];\n group \"Filters\"\n [\n test \"SDSS r pivot wavelength ~ 620 nm\" test_filters_sdss_pivot;\n test \"Johnson V pivot wavelength ~ 551 nm\"\n test_filters_johnson_v_pivot;\n test \"2MASS J pivot wavelength ~ 1235 nm\" test_filters_twomass_j_pivot;\n test \"Gaia BP < G < RP pivot ordering\" test_filters_gaia_ordering;\n test \"5800 K blackbody through SDSS g yields finite magnitude\"\n test_filters_photometry;\n ];\n group \"Survey\"\n [\n test \"Smail n(z) integrates to 1.0\" test_survey_smail_normalized;\n test \"tabulated n(z) is positive at midpoint and zero outside\"\n test_survey_tabulated;\n test \"C_l matrix has correct shape for 2 tracers\" test_survey_cl_shape;\n test \"auto C_l is positive at all ell\" test_survey_cl_positive;\n test \"weak lensing noise is constant in ell\" test_survey_noise_wl;\n test \"shear m_bias=0.02 scales auto C_l by (1+m)^2\" test_survey_m_bias;\n test \"baryonic feedback suppresses high-ell power\"\n test_survey_baryonic_feedback;\n ];\n group \"FITS\"\n [\n test \"2x3 float32 image writes and reads back\"\n test_fits_image_roundtrip;\n test \"3-row table with ra/dec writes and reads back\"\n test_fits_table_roundtrip;\n ];\n group \"Batch\"\n [\n test \"batch of 2 spectra has shape [2; 100]\" test_batch_create;\n test \"mismatched wavelength and values dims raises\"\n test_batch_create_mismatch;\n test \"extract second spectrum from batch\" test_batch_roundtrip;\n test \"scalar scale applies to all spectra in batch\" test_batch_scale;\n test \"batch resample matches per-spectrum resample\"\n test_batch_resample;\n test \"batch AB magnitudes match per-spectrum magnitudes\"\n test_batch_ab_mag;\n test \"batch extinction matches per-spectrum extinction\"\n test_batch_extinction;\n test \"batch redshift matches per-spectrum redshift\"\n test_batch_redshift_scalar;\n ];\n ]\n"
},
{
"path": "doc/coming-from-python.md",
"content": "# Coming from Python\n\nThis page maps Python scientific computing concepts to their Raven equivalents. It assumes you already know OCaml basics.\n\n## Library Mapping\n\n| Python | Raven | Notes |\n|--------|-------|-------|\n| NumPy | [Nx](/docs/nx/) | N-dimensional arrays, broadcasting, linear algebra, FFT |\n| JAX | [Rune](/docs/rune/) | Functional transformations: `grad`, `jvp`, `vmap` |\n| PyTorch / Flax | [Kaun](/docs/kaun/) | Layers, optimizers, training loops |\n| HuggingFace Tokenizers | [Brot](/docs/brot/) | BPE, WordPiece, Unigram; HF-compatible |\n| pandas / Polars | [Talon](/docs/talon/) | Type-safe DataFrames |\n| Matplotlib | [Hugin](/docs/hugin/) | 2D/3D plotting with Cairo |\n| Gymnasium | [Fehu](/docs/fehu/) | RL environments and training utilities |\n| OpenCV | [Sowilo](/docs/sowilo/) | Differentiable image processing |\n| Jupyter + IPython | [Quill](/docs/quill/) | Interactive REPL and markdown notebooks |\n\n## Key Differences\n\n### Explicit Types\n\nNumPy casts types silently. Nx does not.\n\n```python\n# Python: silently upcasts int + float -> float\na = np.array([1, 2, 3])\nb = a + 1.5 # works\n```\n\n\n```ocaml\n(* OCaml: types must match *)\nlet a = Nx.create Nx.Int32 [|3|] [|1l; 2l; 3l|]\n(* Nx.add a (Nx.scalar Nx.Float32 1.5) -- type error *)\n\n(* Cast explicitly *)\nlet a_f = Nx.astype Nx.Float32 a\nlet b = Nx.add a_f (Nx.scalar Nx.Float32 1.5)\n```\n\n### Array Literals\n\nNumPy uses Python lists. Nx uses OCaml arrays with `[| |]` syntax.\n\n```python\nx = np.array([[1, 2], [3, 4]])\n```\n\n\n```ocaml\nlet x = Nx.create Nx.Float32 [|2; 2|] [|1.; 2.; 3.; 4.|]\n```\n\n### Slicing\n\nNumPy uses `[]` with `:`. Nx uses the `slice` function with index constructors.\n\n```python\nx[0:2, :] # first two rows\nx[:, 1] # second column\nx[::2] # every other element\n```\n\n\n```ocaml\nNx.slice [R (0, 2); A] x (* first two rows *)\nNx.slice [A; I 1] x (* second column *)\nNx.slice [S (0, -1, 2)] x (* every other element *)\n```\n\n### No Separate Tensor Type\n\nIn PyTorch, `torch.Tensor` is different from `numpy.ndarray`. In Raven, Rune operates directly on `Nx.t` values. There is no wrapper type.\n\n```python\n# PyTorch: convert between types\nx_np = np.array([1.0, 2.0])\nx_torch = torch.from_numpy(x_np)\nx_torch.requires_grad_(True)\n```\n\n\n```ocaml\n(* Raven: just use Nx tensors directly *)\nlet x = Nx.create Nx.Float32 [|2|] [|1.0; 2.0|]\nlet gradient = Rune.grad (fun x -> Nx.sum (Nx.mul x x)) x\n```\n\n### Functional Transformations\n\nJAX users will find Rune familiar. PyTorch users: think of `grad` as a function transformer, not a method on tensors.\n\n```python\n# JAX style\ngrad_fn = jax.grad(loss_fn)\ngrads = grad_fn(params)\n\n# PyTorch style\nloss = loss_fn(params)\nloss.backward()\ngrads = params.grad\n```\n\n\n```ocaml\n(* Rune: JAX-style functional transforms *)\nlet grad_fn = Rune.grad loss_fn\nlet grads = grad_fn params\n\n(* Or compute value and gradient together *)\nlet loss, grads = Rune.value_and_grad loss_fn params\n```\n\n### Module-Based Layers\n\nKaun layers are records with `init` and `apply`, not classes with `forward`.\n\n```python\n# PyTorch\nclass Model(nn.Module):\n def __init__(self):\n self.linear = nn.Linear(784, 10)\n def forward(self, x):\n return self.linear(x)\nmodel = Model()\n```\n\n\n```ocaml\n(* Kaun: compose layer records *)\nlet model = Kaun.Layer.sequential [\n Kaun.Layer.linear ~in_features:784 ~out_features:10 ();\n]\nlet vars = Kaun.Layer.init model ~dtype:Nx.Float32\n```\n\nParameters are plain data (`Ptree.t` — a tree of Nx tensors), not hidden inside objects.\n\n### DataFrames\n\npandas uses string-based column access. Talon provides type-safe row operations via an applicative.\n\n```python\n# pandas\ndf['bmi'] = df['weight'] / df['height'] ** 2\n```\n\n\n```ocaml\n(* Talon: type-safe row computation *)\nlet df = Talon.with_column df \"bmi\" Nx.Float64\n Talon.Row.(map2 (number \"weight\") (number \"height\")\n ~f:(fun w h -> w /. (h *. h)))\n```\n\n## Detailed Comparisons\n\nEach library has a dedicated comparison page with side-by-side code examples:\n\n- [Nx vs NumPy](/docs/nx/numpy-comparison/)\n- [Rune vs JAX](/docs/rune/jax-comparison/)\n- [Kaun vs PyTorch/Flax](/docs/kaun/pytorch-comparison/)\n- [Brot vs HuggingFace Tokenizers](/docs/brot/hf-tokenizers-comparison/)\n- [Talon vs pandas](/docs/talon/pandas-comparison/)\n- [Hugin vs Matplotlib](/docs/hugin/matplotlib-comparison/)\n- [Sowilo vs OpenCV](/docs/sowilo/opencv-comparison/)\n- [Fehu vs Gymnasium](/docs/fehu/gymnasium-comparison/)\n"
},
{
"path": "doc/ecosystem-overview.md",
"content": "# The Raven Ecosystem\n\nRaven is nine libraries that share one data type: `Nx.t`, the\nn-dimensional array. Each library does one thing, and they compose\nthrough tensors.\n\n## How the Libraries Fit Together\n\n```\n ┌───────────┐\n │ Kaun │ neural networks\n │ (Flax) │\n └─────┬─────┘\n │\n ┌───────────┐ ┌─────┴─────┐ ┌───────────┐\n │ Sowilo │ │ Rune │ │ Fehu │\n │ (OpenCV) ├──────────┤ (JAX) ├──────────┤(Gymnasium)│\n └─────┬─────┘ └─────┬─────┘ └─────┬─────┘\n │ │ │\n ┌─────┴──────────────────────┴──────────────────────┴─────┐\n │ Nx │\n │ (NumPy) │\n └──┬──────────────┬──────────────┬──────────────┬─────────┘\n │ │ │ │\n ┌───┴────┐ ┌────┴───┐ ┌────┴───┐ ┌─────┴────┐\n │ Talon │ │ Brot │ │ Hugin │ │ Quill │\n │(Polars)│ │(HF Tok)│ │(Mpl) │ │(Jupyter) │\n └────────┘ └────────┘ └────────┘ └──────────┘\n```\n\n**Nx** is the foundation — every library operates on `Nx.t` tensors.\n\n**Rune** adds functional transformations on top of Nx: `grad`, `jvp`,\n`vmap`. Your Nx code becomes differentiable without changes.\n\n**Kaun** builds on Rune to provide layers, optimizers, training loops,\nand HuggingFace Hub integration.\n\n**Sowilo**, **Fehu**, **Talon**, **Brot**, **Hugin**, and **Quill** each\nuse Nx directly for their domain. Sowilo and Fehu operations are\ncompatible with Rune's `grad` and `vmap` since they are plain Nx\noperations under the hood.\n\n## Which Library Do I Need?\n\n| I want to... | Use |\n|---|---|\n| Work with numerical arrays | [Nx](/docs/nx/) |\n| Compute gradients | [Rune](/docs/rune/) |\n| Train neural networks | [Kaun](/docs/kaun/) |\n| Tokenize text for language models | [Brot](/docs/brot/) |\n| Manipulate tabular data | [Talon](/docs/talon/) |\n| Process and transform images | [Sowilo](/docs/sowilo/) |\n| Build RL environments and agents | [Fehu](/docs/fehu/) |\n| Create plots and visualizations | [Hugin](/docs/hugin/) |\n| Run code interactively (REPL or notebooks) | [Quill](/docs/quill/) |\n\n---\n\n## Nx: N-Dimensional Arrays\n\nNx provides the numerical foundation for the entire ecosystem.\nNumPy-like operations on n-dimensional arrays with 19 data types\n(float16 through complex128), broadcasting, slicing, linear algebra,\nFFT, and I/O.\n\n```ocaml\nopen Nx\n\nlet x = linspace Float32 0. 10. 100\nlet y = sin x\nlet mean_y = mean y\n```\n\n[Nx documentation →](/docs/nx/)\n\n## Rune: Automatic Differentiation\n\nFunctional transformations for Nx tensors: reverse-mode AD (grad,\nvjp), forward-mode AD (jvp), and vectorising maps (vmap). Operates on\n`Nx.t` values directly using OCaml 5 effect handlers — no special\ntensor type needed.\n\n\n```ocaml\nopen Nx\nopen Rune\n\nlet f x = add (mul x x) (sin x)\nlet f' = grad f\nlet f'' = grad f'\n```\n\n[Rune documentation →](/docs/rune/)\n\n## Kaun: Neural Networks\n\nComposable layers, optimizers with learning-rate schedules, training\nloops, data pipelines, and HuggingFace Hub integration. Model\nparameters are `Ptree.t` — trees of Nx tensors you can inspect, map,\nand serialize.\n\n\n```ocaml\nopen Kaun\n\nlet model = Layer.sequential [\n Layer.linear ~in_features:784 ~out_features:128 ();\n Layer.relu ();\n Layer.linear ~in_features:128 ~out_features:10 ();\n]\n\nlet trainer = Train.make ~model\n ~optimizer:(Optim.adam ~lr:(Optim.Schedule.constant 0.001) ())\n```\n\n[Kaun documentation →](/docs/kaun/)\n\n## Brot: Tokenization\n\nFast, HuggingFace-compatible tokenization supporting BPE, WordPiece,\nUnigram, word-level, and character-level algorithms. Composable\npipeline (normalizer → pre-tokenizer → model → post-processor →\ndecoder) with training from scratch.\n\n```ocaml\nopen Brot\n\nlet tokenizer = from_file \"tokenizer.json\" |> Result.get_ok\nlet encoding = encode tokenizer \"Hello, world!\"\nlet ids = Encoding.ids encoding\n```\n\n[Brot documentation →](/docs/brot/)\n\n## Talon: DataFrames\n\nType-safe tabular data with heterogeneous columns, an applicative Row\nsystem for row-wise operations, and vectorized aggregations backed by\nNx.\n\n```ocaml\nopen Talon\n\nlet df = create [\n \"name\", Col.string_list [\"Alice\"; \"Bob\"; \"Charlie\"];\n \"score\", Col.float64_list [85.5; 92.0; 78.5];\n]\n\nlet () = print df\n```\n\n[Talon documentation →](/docs/talon/)\n\n## Sowilo: Computer Vision\n\nDifferentiable image processing: geometric transforms (resize, crop,\nflip), spatial filters (Gaussian blur, Sobel, Canny), color space\nconversions, and morphological operations. All operations are plain Nx\ncomputations, so they compose with `Rune.grad` and `Rune.vmap`.\n\n\n```ocaml\nopen Sowilo\n\nlet processed =\n img\n |> to_float\n |> resize ~height:224 ~width:224 ~mode:Bilinear\n |> normalize ~mean:[|0.485; 0.456; 0.406|] ~std:[|0.229; 0.224; 0.225|]\n```\n\n[Sowilo documentation →](/docs/sowilo/)\n\n## Fehu: Reinforcement Learning\n\nRL environments (CartPole, MountainCar, GridWorld), type-safe\nobservation/action spaces, vectorized environments, trajectory\ncollection, replay buffers, and generalized advantage estimation.\n\n\n```ocaml\nopen Fehu\n\nlet env = Fehu_envs.cartpole () in\nlet obs, _info = Env.reset env in\nlet obs, reward, terminated, truncated, _info =\n Env.step env (Space.sample (Env.action_space env))\n```\n\n[Fehu documentation →](/docs/fehu/)\n\n## Hugin: Visualization\n\nPublication-quality 2D and 3D plots using Cairo rendering. Takes Nx\ntensors as input. Line plots, scatter, bar charts, contour plots,\nimage display.\n\n\n```ocaml\nopen Hugin\nopen Nx\n\nlet fig = figure () in\nlet ax = subplot fig in\nlet _ = Plotting.plot ax ~x ~y ~label:\"sin(x)\" in\nshow fig\n```\n\n[Hugin documentation →](/docs/hugin/)\n\n## Quill: Interactive Computing\n\nInteractive REPL and markdown notebooks. Launch `quill` for a toplevel\nwith syntax highlighting, completion, and history, or open a markdown\nfile for a full notebook experience. Terminal UI, web frontend, and\nbatch mode with all Raven libraries pre-loaded.\n\n\n```bash\nquill # interactive REPL\nquill notebook.md # notebook TUI\nquill serve notebook.md # web frontend\nquill run notebook.md # batch evaluation\n```\n\n[Quill documentation →](/docs/quill/)\n\n## Getting Started\n\n1. **New to Raven?** Start with the [Quickstart](/docs/quickstart/)\n2. **Coming from Python?** Read [Coming from Python](/docs/coming-from-python/)\n3. **Want a specific library?** Use the table above to find the right docs\n"
},
{
"path": "doc/index.md",
"content": "# Documentation\n\nWelcome to Raven's documentation. Raven is an ecosystem of OCaml libraries for numerical computing, machine learning, and data science.\n\n## Start Here\n\n- **[Quickstart](/docs/quickstart/)** — zero to gradient in 5 minutes\n- **[Coming from Python](/docs/coming-from-python/)** — map NumPy, PyTorch, pandas concepts to Raven\n- **[Ecosystem Overview](/docs/ecosystem-overview/)** — how the libraries relate and which to use\n\n## Libraries\n\n| | Library | Like | What it does |\n|-|---------|------|-------------|\n| | [**nx**](/docs/nx/) | NumPy | N-dimensional arrays with pluggable backends |\n| ᚱ | [**rune**](/docs/rune/) | JAX | Automatic differentiation and functional transformations |\n| ᚲ | [**kaun**](/docs/kaun/) | PyTorch / Flax | Neural networks and training |\n| ᚨ | [**brot**](/docs/brot/) | HF Tokenizers | Fast tokenization for language models |\n| ᛃ | [**talon**](/docs/talon/) | Pandas / Polars | DataFrames with type-safe columns |\n| ᛞ | [**hugin**](/docs/hugin/) | Matplotlib | Data visualization and plotting |\n| ᛈ | [**quill**](/docs/quill/) | Jupyter + IPython | Interactive REPL and markdown notebooks |\n| ᚠ | [**fehu**](/docs/fehu/) | Gymnasium | Reinforcement learning environments |\n| ᛋ | [**sowilo**](/docs/sowilo/) | OpenCV | Differentiable computer vision |\n\n## Project\n\n- [Installation](/docs/installation/) — system dependencies, opam setup, building from source\n- [Roadmap](/docs/roadmap/) — what works today and what's coming\n- [Introduction](/docs/introduction/) — vision and philosophy\n- [Support Raven](/docs/support-raven/) — sponsorship and contributing\n"
},
{
"path": "doc/installation.md",
"content": "# Installation\n\n## Prerequisites\n\nRaven requires **OCaml 5.2** or later and **opam**.\n\nIf you don't have opam installed, follow the [official instructions](https://opam.ocaml.org/doc/Install.html). Then create a switch:\n\n```bash\nopam switch create raven 5.2.0\neval $(opam env)\n```\n\n## Installing from opam\n\nInstall the entire ecosystem:\n\n```bash\nopam install raven\n```\n\nOr install individual libraries:\n\n```bash\nopam install nx # just arrays\nopam install rune # arrays + autodiff\nopam install kaun # arrays + autodiff + neural networks\nopam install brot # tokenization\nopam install talon # dataframes\n```\n\n## Building from Source\n\n```bash\ngit clone https://github.com/raven-ml/raven\ncd raven\ndune pkg lock && dune build\n```\n\nTo build a specific library:\n\n```bash\ndune build packages/nx # just nx\ndune build packages/kaun # kaun + its dependencies\n```\n\n## System Dependencies\n\nMost Raven libraries have no system dependencies beyond OCaml. The exceptions:\n\n| Library | Requires | macOS | Ubuntu/Debian |\n|---------|----------|-------|---------------|\n| **hugin** | Cairo, SDL2 | `brew install cairo sdl2` | `apt install libcairo2-dev libsdl2-dev` |\n\n## Using Raven in Your Project\n\nAdd libraries to your `dune-project`:\n\n```dune\n(lang dune 3.0)\n(package\n (name my_project)\n (depends\n ocaml\n dune\n nx\n rune))\n```\n\nAnd your `dune` file:\n\n```dune\n(executable\n (name main)\n (libraries nx rune))\n```\n\n## Verify Your Installation\n\nCreate a file `main.ml`:\n\n```ocaml\nlet () =\n let open Nx in\n let x = linspace Float32 0. 1. 5 in\n print_data x\n```\n\nBuild and run:\n\n```bash\ndune exec ./main.exe\n```\n\nYou should see five evenly-spaced values printed.\n\n## Editor Setup\n\nFor the best development experience, use an editor with OCaml LSP support:\n\n- **VS Code**: Install the [OCaml Platform extension](https://marketplace.visualstudio.com/items?itemName=ocamllabs.ocaml-platform)\n- **Emacs**: Use [ocaml-eglot](https://github.com/tarides/ocaml-eglot)\n- **Vim/Neovim**: Use [ocaml-lsp](https://github.com/ocaml/ocaml-lsp) with your LSP client\n\n## Troubleshooting\n\n**Missing system libraries**: If Hugin fails to build, ensure Cairo and SDL2 development headers are installed.\n\n**Opam switch issues**: Run `eval $(opam env)` after creating or switching opam switches.\n\n**Build failures**: Check your OCaml version with `ocaml --version`. Raven requires 5.2.0 or later.\n\n**Getting help**: Report issues at [github.com/raven-ml/raven/issues](https://github.com/raven-ml/raven/issues).\n"
},
{
"path": "doc/introduction.md",
"content": "# Introduction\n\nRaven is a project to bring modern scientific computing to the OCaml programming language. We're building a comprehensive ecosystem, from low-level numerical libraries and automatic differentiation to high-level machine learning frameworks and interactive notebooks.\n\nOur ambition is to make scientific computing in OCaml feel as natural as it does in Python. This means not just matching Python's capabilities, but delivering the same level of ergonomics, performance, and developer experience that has made Python the de facto standard for scientific computing.\n\nIf successful, Raven would establish OCaml as a genuine alternative in the scientific computing landscape. It's an ambitious undertaking, but one we believe is both necessary and achievable.\n\n## Why Not Just Use Python?\n\nToday, Python has an effective monopoly on scientific computing. Unlike web development, where we can choose between multiple mature ecosystems, numerical computing offers essentially one realistic option. This lack of choice is unfortunate.\n\nWhat's more problematic is that Python, while excellent for quick experimentation, doesn't particularly shine for building robust production systems. Its interpreted nature, dynamic typing, and limited multicore support create real challenges when you need to deploy and maintain large-scale applications.\n\nIf you've worked in this space, you've likely experienced this firsthand: rapid prototypes that become production nightmares, debugging sessions where type errors only surface at runtime, or performance bottlenecks that force you to drop down to C extensions.\n\nOften, this mismatch forces a wasteful pattern: researchers prototype in Python, then teams reimplement everything for production in other languages. This induces all kinds of second-order effects on organization structures, team dynamics, development velocity, and workload.\n\nThe scientific community deserves better options than being forced into one language, and we believe OCaml occupies a unique sweet spot between rapid experimentation and building production-grade systems. It just needs the scientific ecosystem to match its technical strengths. This is the gap that Raven aims to fill.\n\nIn the AI era, we believe OCaml has an important role to play. If you're generating 80% of your code with AI assistance, wouldn't you prefer a language that catches errors at compile time rather than runtime? The productivity gains from AI coding are amplified when you have a type system that gives you stronger guarantees about your generated code. Raven is our contribution to putting OCaml in the spotlight for scientific computing in this new era.\n\n## What Does Success Look Like?\n\nOur goal isn't just to build OCaml versions of Python libraries: it's to create a compelling alternative for busy developers who just want the best tool for the job.\n\nSuccess means two things. First, **OCaml developers shouldn't have to switch to Python for numerical computing**. Whether you're analyzing data, training models, or building computational systems, you should be able to stay in the OCaml ecosystem with the same productivity you'd expect from Python.\n\nSecond, **Raven should break into the mainstream scientific computing conversation**. It shouldn't just serve existing OCaml developers: we're building for teams who need to ship reliable systems, not just an OCaml curiosity for language enthusiasts.\n\nWe measure success across five key dimensions:\n\n- **Capability parity**: Everything you can do in Python, you should be able to do with Raven\n- **Development productivity**: Getting from idea to working prototype is as fast as it would be in Python\n- **Developer experience**: Developers get the kind of documentation, tooling, and APIs they dream every project had\n- **Production performance**: Match or exceed NumPy/PyTorch performance on the fast path\n- **Production readiness**: Teams can ship robust, maintainable Raven-built applications that perform well under real-world conditions\n\nWe believe this is achievable through focused execution and strategic choices. We're prioritizing the 80% that matter most, focusing on one blessed workflow per use-case, and building modular components that encourage ecosystem growth, rather than trying to match Python everywhere from day one.\n\n## Why Not Just Use Owl?\n\nOwl deserves credit for the amount of work and love that has been poured into it. It demonstrated that serious numerical computing in OCaml was possible, spanning everything from statistics and signal processing to basic linear algebra and neural networks, and more.\n\nHowever, Owl can't compete with NumPy or PyTorch on performance, and performance parity isn't optional if we want teams to seriously consider OCaml over Python.\n\nThe reality is that we can't realistically match NumPy and PyTorch's performance through traditional optimization. These projects have hundreds of developers working on hand-optimized kernels. With our small team, JIT compilation is our only viable path to competitive performance.\n\nThis creates a fundamental constraint. Building for JIT-first changes everything about your design: API choices, memory layouts, operator fusion strategies, even how you structure the development experience. Rather than retrofitting these assumptions onto existing work, we decided a clean slate would be more effective.\n\nThere's also the ecosystem question. Despite Owl's technical achievements, it hasn't generated the kind of flourishing community we need. We suspect this is partly due to its lack of modularity: without libraries designed as composable building blocks, it's challenging to build a broader ecosystem around the foundation.\n\nRaven is designed from the ground up to (1) compete with Python's scientific computing stack on performance and (2) build the flourishing ecosystem that OCaml's scientific computing community deserves.\n\n## What We're Building\n\nRaven is a comprehensive ecosystem that spans the entire scientific computing stack. Here's what we're building:\n\n**Foundation**\n- **Nx**: N-dimensional arrays with pluggable backends (NumPy equivalent)\n- **Brot**: Fast, HuggingFace-compatible tokenization (HF Tokenizers equivalent)\n- **Talon**: Type-safe DataFrames (pandas/Polars equivalent)\n\n**Differentiable Computing**\n- **Rune**: Automatic differentiation using OCaml's effect system (JAX equivalent)\n\n**Domain Frameworks**\n- **Kaun**: Neural networks and training (PyTorch/Flax equivalent)\n- **Sowilo**: Differentiable computer vision (OpenCV equivalent)\n- **Fehu**: Reinforcement learning environments and algorithms (Gymnasium equivalent)\n\n**Tooling**\n- **Hugin**: Publication-quality plotting (Matplotlib equivalent)\n- **Quill**: Interactive notebooks as markdown files (Jupyter equivalent)\n\nNine libraries spanning the full scientific computing stack, all designed to work together seamlessly.\n\n**Key Innovations**\nWhile we aim to feel familiar to Python users, Raven brings genuine innovations to scientific computing:\n\n**Nx** uses pluggable backends inspired by Tinygrad's minimalist approach, giving us flexibility to optimize for different hardware without monolithic complexity.\n\n**Rune** implements automatic differentiation using OCaml's effects system. As far as we know, it is the first project of this scale to use effects for autodiff, building on recent research, and implementing Jax's vision for functional numerical computation with a truly functional foundation.\n\n**Quill** rethinks notebooks. Notebooks are plain markdown files — git-friendly, readable without special tooling, and editable in any text editor. Quill runs them as a TUI in the terminal or as a web frontend in the browser, with all Raven packages pre-loaded and zero setup.\n\n**Deployment** is where Raven's story diverges most from Python. AOT compilation generates all compute kernels at compile time, producing binaries with no BLAS or CUDA runtime dependency. This makes it possible to deploy models as MirageOS unikernels — minimal attack surface, millisecond boot, deterministic behavior — or as static binaries with no Python runtime, no dependency hell.\n\n**Current Focus**\nThe alpha milestone is complete — we've trained GPT-2 end-to-end on CPU using the full Raven stack. We're now focused on integrating tolk as a JIT transformation in Rune, with the goal of matching PyTorch performance. After that, V1 brings production-ready training and deployment: AOT compilation, inference serving, ONNX import, and MirageOS unikernel deployment. See the [roadmap](/docs/roadmap/) for details.\n"
},
{
"path": "doc/quickstart.md",
"content": "# Quickstart\n\nThis gets you from zero to computing gradients and training a model in five minutes.\n\n## Setup\n\n\n```bash\nopam install raven\n```\n\nCreate a `dune-project` and `dune` file:\n\n\n```dune\n; dune-project\n(lang dune 3.20)\n```\n\n\n```dune\n; dune\n(executable\n (name main)\n (libraries kaun))\n```\n\nInstalling `kaun` pulls in `nx` and `rune` automatically.\n\n## Step 1: Arrays with Nx\n\nNx provides n-dimensional arrays. Every value has a data type and a shape.\n\n```ocaml\nopen Nx\n\nlet () =\n (* Create arrays *)\n let a = create Float32 [|2; 3|] [|1.; 2.; 3.; 4.; 5.; 6.|] in\n let b = ones Float32 [|2; 3|] in\n\n (* Element-wise operations *)\n let c = add a b in\n print_data c;\n\n (* Reductions *)\n Printf.printf \"sum = %.1f\\n\" (item [] (sum a));\n Printf.printf \"mean = %.1f\\n\" (item [] (mean a));\n\n (* Matrix multiplication *)\n let x = rand Float32 [|3; 4|] in\n let y = rand Float32 [|4; 2|] in\n let z = matmul x y in\n Printf.printf \"matmul shape: %s\\n\"\n (Array.to_list (shape z) |> List.map string_of_int |> String.concat \"x\")\n```\n\n## Step 2: Gradients with Rune\n\nRune computes derivatives of Nx functions automatically. Write a function using Nx operations, then use `grad` to differentiate it.\n\n```ocaml\nopen Nx\nopen Rune\n\nlet () =\n (* f(x) = x² + sin(x) *)\n let f x = add (mul x x) (sin x) in\n\n (* grad returns the derivative function *)\n let f' = grad f in\n\n let x = scalar Float32 2.0 in\n Printf.printf \"f(2) = %.4f\\n\" (item [] (f x));\n Printf.printf \"f'(2) = %.4f\\n\" (item [] (f' x));\n\n (* Higher-order: second derivative *)\n let f'' = grad f' in\n Printf.printf \"f''(2) = %.4f\\n\" (item [] (f'' x))\n```\n\n## Step 3: Training with Kaun\n\nKaun provides layers, optimizers, and training loops built on Rune.\n\n\n```ocaml\nopen Kaun\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n\n (* XOR dataset *)\n let x = Nx.create Nx.Float32 [|4; 2|]\n [|0.; 0.; 0.; 1.; 1.; 0.; 1.; 1.|] in\n let y = Nx.create Nx.Float32 [|4; 1|]\n [|0.; 1.; 1.; 0.|] in\n\n (* Define model *)\n let model = Layer.sequential [\n Layer.linear ~in_features:2 ~out_features:8 ();\n Layer.tanh ();\n Layer.linear ~in_features:8 ~out_features:1 ();\n ] in\n\n (* Create trainer and initialize *)\n let trainer = Train.make ~model\n ~optimizer:(Optim.adam ~lr:(Optim.Schedule.constant 0.01) ()) in\n let st = Train.init trainer ~dtype:Nx.Float32 in\n\n (* Train *)\n let st = Train.fit trainer st\n ~report:(fun ~step ~loss _st ->\n if step mod 250 = 0 then\n Printf.printf \"step %4d loss %.6f\\n\" step loss)\n (Data.repeat 1000 (x, fun pred -> Loss.binary_cross_entropy pred y))\n in\n\n (* Predict *)\n let pred = Train.predict trainer st x |> Nx.sigmoid in\n Printf.printf \"\\npredictions (expected 0 1 1 0):\\n\";\n for i = 0 to 3 do\n Printf.printf \" [%.0f, %.0f] -> %.3f\\n\"\n (Nx.item [i; 0] x) (Nx.item [i; 1] x) (Nx.item [i; 0] pred)\n done\n```\n\n## Next Steps\n\n- **[Nx](/docs/nx/getting-started/)** — full guide to arrays, slicing, broadcasting, linear algebra\n- **[Rune](/docs/rune/getting-started/)** — all transformations: grad, jvp, vmap, and more\n- **[Kaun](/docs/kaun/getting-started/)** — layers, optimizers, training loops, pretrained models\n- **[Ecosystem Overview](/docs/ecosystem-overview/)** — how all 9 libraries fit together\n"
},
{
"path": "doc/roadmap.md",
"content": "# Roadmap\n\n## Current Status\n\nRaven is in **alpha**. The core stack (Nx -> Rune -> Kaun) works end-to-end: we have successfully trained GPT-2 on CPU using the full Raven stack.\n\n| Library | Status | What works |\n| ---------- | ------ | ------------------------------------------------------------------------- |\n| **nx** | Alpha | Full NumPy-like API, linear algebra, FFT, I/O (npy, images) |\n| **rune** | Alpha | Reverse and forward-mode AD, vmap, gradient checking |\n| **kaun** | Alpha | Layers, optimizers, training loops, HuggingFace Hub, MNIST/GPT-2 examples |\n| **brot** | Alpha | All 5 algorithms, full pipeline, HF tokenizer.json compat, training |\n| **talon** | Alpha | DataFrames, row operations, aggregations, CSV I/O |\n| **hugin** | Alpha | 2D/3D plots, scatter, bar, contour, images |\n| **fehu** | Alpha | Environments (CartPole, GridWorld, MountainCar), vectorized envs, GAE |\n| **sowilo** | Alpha | Geometric transforms, filters, edge detection, morphological ops |\n| **quill** | Alpha | Interactive REPL, notebook TUI and web frontend, batch eval, watch mode |\n\nAPIs will change. Bug reports and feedback are welcome.\n\n## Beta: JIT Compilation & Performance\n\nThe beta cycle focuses on **JIT compilation with performance close to PyTorch**.\n\n- Integrate tolk (an OCaml port of tinygrad) as a JIT transformation in Rune\n- Target CPU, CUDA, Metal, OpenCL, and HIP\n- Kernel fusion and optimization\n- Benchmark against PyTorch on standard workloads\n\n## V1: Production-Ready Training & Deployment\n\nV1 makes Raven **production-ready**: train models, deploy them as unikernels or static binaries.\n\n**Training**:\n- Gradient accumulation, mixed precision, gradient checkpointing\n- Flash attention for efficient transformer training\n- ONNX import for PyTorch model portability\n- Parallel data loading, layer completions\n\n**Deployment**:\n- AOT compilation to standalone binaries (CPU and GPU)\n- Inference engine with KV cache, continuous batching, and PagedAttention\n- Post-training quantization (INT8/INT4)\n- MirageOS unikernel deployment -- tolk AOT generates all compute at compile time, no BLAS dependency, enabling deployment as unikernels\n"
},
{
"path": "doc/support-raven.md",
"content": "# Support Raven\n\n## Raven in One Minute\n\nPython's monopoly on scientific computing forces an impossible choice: ship everything in Python (endure runtime crashes, the GIL's multicore ceiling, and gigabyte containers), or prototype in Python then rewrite for production (doubling the work and creating siloed teams).\n\n**We think there's a better way.** OCaml lets you prototype as quickly as Python and scale the same code to production. Same expressiveness, strong typing catches bugs before they crash your ML pipeline, while JIT compilation matches NumPy/PyTorch performance. One language from research to production — it just needs a production-grade ML stack.\n\n**Raven brings that stack to OCaml:** Nx (NumPy), Rune (JAX with effects-based autodiff), Kaun (Flax), Brot (tokenization), Hugin (Matplotlib), and Quill (notebooks done right). Train models with automatic differentiation and JIT compilation, then deploy as a MirageOS unikernel or a static binary — no Python, no CUDA dependency hell, no 5 GB Docker images. We built Raven for teams that want both development speed and reliable systems.\n\n_Learn more: [Introduction](/docs/introduction)_\n\n_We're in alpha with the full stack working end-to-end (we've trained GPT-2 on CPU). Next milestone: JIT compilation via tolk with performance close to PyTorch._\n\n## Roadmap & Funding Goals\n\n_See the [full roadmap](/docs/roadmap) for our complete vision and timeline._\n\n### Beta — JIT Compilation & Performance\n- Integrate tolk (tinygrad-based compiler) as a JIT transformation in Rune\n- Target CPU, CUDA, Metal, OpenCL, and HIP\n- Kernel fusion and optimization\n- Performance within 2x of PyTorch on standard workloads\n\n### V1 — Production-Ready Training & Deployment\n- Production training: gradient accumulation, mixed precision, gradient checkpointing, flash attention\n- ONNX import for PyTorch model portability\n- AOT compilation to standalone binaries (CPU and GPU)\n- Inference engine with KV cache, continuous batching, and PagedAttention\n- MirageOS unikernel deployment\n- Post-training quantization (INT8/INT4)\n\nWe're also open to discussing custom sponsorship packages based on your needs.\n\n## Ways to Support\n\n### For Developers\n- **Try it out**: Test Raven with your workflows and [report issues](https://github.com/raven-ml/raven/issues)\n- **Contribute code**: See our [contributing guide](https://github.com/raven-ml/raven/blob/main/CONTRIBUTING.md) for areas where we need help\n- **Share feedback**: What would make you switch from Python? [Tell us](mailto:thibaut.mattio@gmail.com)\n- **Spread the word**: Star the repo, share with your team, write about your experience\n\n### For Companies\n- **Use Raven**: Reach out if you're interested in using it—we're keen on prioritizing development based on real-world needs\n- **Sponsor development**: Email [thibaut.mattio@gmail.com](mailto:thibaut.mattio@gmail.com) for sponsorship packages\n\n### For Individuals\n- **GitHub Sponsors**: [Support the project with monthly contributions](https://github.com/sponsors/tmattio)\n- **One-time donations**: Every contribution helps us reach the next milestone\n- **Write tutorials**: Help others learn Raven and grow the community\n\n## Current Sponsors\n\nWe're grateful for the support of our sponsors:\n\n### Corporate Sponsors\n\n- [**Ahrefs**](https://ahrefs.com) - Building tools to help you grow your search traffic\n- [**Tarides**](https://tarides.com) - Secure-by-design infrastructure and tooling for a better digital world\n\n### Individual Sponsors\n\nThank you to all our individual sponsors for their support!\n\n## Get in Touch\n\n**For sponsorship inquiries**: [thibaut.mattio@gmail.com](mailto:thibaut.mattio@gmail.com) \n**For feature request or bug reports**: [GitHub Issues](https://github.com/raven-ml/raven/issues)\n\n---\n\n_Raven is built by [Thibaut Mattio](https://github.com/tmattio) and contributors. We believe OCaml deserves a world-class scientific computing ecosystem, and we're committed to building it._\n"
},
{
"path": "dune-project",
"content": "(lang dune 3.21)\n\n(name raven)\n\n(source\n (github raven-ml/raven))\n\n(authors \"Thibaut Mattio \")\n\n(maintainers \"Thibaut Mattio \")\n\n(license ISC)\n\n(documentation \"https://raven-ml.dev/docs/\")\n\n(bug_reports \"https://github.com/raven-ml/raven/issues\")\n\n(using mdx 0.4)\n\n(using directory-targets 0.1)\n\n(version 1.0.0~alpha3)\n\n(implicit_transitive_deps false)\n\n(generate_opam_files true)\n\n(opam_file_location inside_opam_directory)\n\n(pin\n (url \"git+https://github.com/invariant-hq/thumper.git\")\n (package\n (name thumper)))\n\n(package\n (name nx)\n (dir packages/nx)\n (synopsis \"N-dimensional arrays for OCaml\")\n (description\n \"Nx provides n-dimensional arrays with NumPy-like semantics and OCaml's type safety. 19 data types, broadcasting, slicing, linear algebra, FFT, and I/O. The numerical foundation for the Raven ecosystem.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (dune-configurator :build)\n (conf-pkg-config :build)\n ; camlzip\n (conf-zlib :build)\n logs\n ; tests\n (windtrap :with-test)\n (mdx :with-test)\n (thumper :with-test))\n (tags\n (numerical-computation tensor-library machine-learning)))\n\n(package\n (name brot)\n (dir packages/brot)\n (synopsis \"Tokenization for OCaml\")\n (description\n \"Fast, HuggingFace-compatible tokenization for language models. BPE, WordPiece, Unigram, word-level, and character-level algorithms with composable pipelines and training from scratch.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n re\n jsont\n bytesrw\n (uunf\n (>= 15.1.0))\n uucp\n (windtrap :with-test)\n (mdx :with-test)\n (thumper :with-test))\n (tags\n (tokenization bpe wordpiece subword-tokenization language-models)))\n\n(package\n (name talon)\n (dir packages/talon)\n (synopsis \"Dataframes for OCaml\")\n (description\n \"Fast and elegant dataframes with type-safe operations. Heterogeneous columns, applicative row operations, vectorized aggregations, and CSV I/O, built on Nx.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (nx\n (= :version))\n (windtrap :with-test)\n (mdx :with-test)\n (thumper :with-test))\n (tags\n (dataframe data-manipulation data-science tabular-data)))\n\n(package\n (name rune)\n (dir packages/rune)\n (synopsis \"Functional transformations for Nx arrays\")\n (description\n \"Automatic differentiation and vectorizing maps for Nx tensors. Reverse-mode AD (grad, vjp), forward-mode AD (jvp), vmap, and gradient checking, built on OCaml 5 effect handlers.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (dune-configurator :build)\n (nx\n (= :version))\n (tolk\n (= :version))\n (windtrap :with-test)\n (mdx :with-test)\n (thumper :with-test))\n (tags\n (automatic-differentiation machine-learning deep-learning optimization)))\n\n(package\n (name tolk)\n (dir packages/tolk)\n (synopsis \"A minimal ML compiler for GPU tensor computation\")\n (description\n \"Tolk is a minimal, readable ML compiler for GPU tensor computation in the Raven ecosystem.\")\n (depends\n (ocaml\n (>= 5.2))\n dune\n (windtrap :with-test)\n (thumper :with-test))\n (tags\n (compiler gpu tensor-computation)))\n\n(package\n (name norn)\n (dir packages/norn)\n (synopsis \"MCMC sampling for OCaml\")\n (description\n \"Markov chain Monte Carlo samplers with automatic gradients via Rune. Hamiltonian Monte Carlo with dual-averaging step-size adaptation.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (nx\n (= :version))\n (rune\n (= :version))\n (windtrap :with-test)\n (thumper :with-test))\n (tags\n (mcmc sampling bayesian machine-learning)))\n\n(package\n (name vega)\n (dir packages/vega)\n (synopsis \"Per-parameter gradient-based optimizers for OCaml\")\n (description\n \"Typed, per-parameter optimizer primitives: Adam, AdamW, SGD, RMSprop, Adagrad, and learning-rate schedules. Built on Nx with no autodiff dependency.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (nx\n (= :version))\n (windtrap :with-test)\n (thumper :with-test))\n (tags\n (optimization machine-learning gradient-descent)))\n\n(package\n (name kaun)\n (dir packages/kaun)\n (synopsis \"Neural networks for OCaml\")\n (description\n \"Composable layers, parameter trees, optimizers, training loops, data pipelines, and HuggingFace Hub integration. Built on Rune.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (rune\n (= :version))\n (vega\n (= :version))\n (nx\n (= :version))\n jsont\n bytesrw\n (windtrap :with-test)\n (mdx :with-test)\n (thumper :with-test))\n (tags\n (neural-networks machine-learning deep-learning)))\n\n(package\n (name munin)\n (dir packages/munin)\n (synopsis \"Local experiment tracking for Raven\")\n (description\n \"Local-first experiment tracking with append-only event logs, versioned artifacts, a terminal dashboard, and a CLI. The core library (munin) provides Session, Run, Store, and Artifact modules. The TUI library (munin.tui) provides a Mosaic-based dashboard.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n jsont\n bytesrw\n sha\n cmdliner\n mosaic\n (dune-configurator :build)\n matrix\n (windtrap :with-test))\n (tags\n (experiment-tracking machine-learning monitoring)))\n\n(package\n (name sowilo)\n (dir packages/sowilo)\n (synopsis \"Differentiable computer vision for OCaml\")\n (description\n \"Image processing operations expressed as Nx tensor computations. Geometric transforms, spatial filters, edge detection, morphological operations, and color space conversions, all compatible with Rune.grad and Rune.vmap.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (nx\n (= :version))\n (windtrap :with-test)\n (mdx :with-test)\n (thumper :with-test))\n (tags\n (computer-vision image-processing feature-detection machine-learning)))\n\n(package\n (name fehu)\n (dir packages/fehu)\n (synopsis \"Reinforcement learning for OCaml\")\n (description\n \"Type-safe RL environments, observation/action spaces, vectorized environments, trajectory collection, replay buffers, and generalized advantage estimation. Built on Nx.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (nx\n (= :version))\n (windtrap :with-test)\n (mdx :with-test)\n (thumper :with-test))\n (tags\n (reinforcement-learning machine-learning environments)))\n\n(package\n (name hugin)\n (dir packages/hugin)\n (synopsis \"Declarative plotting and visualization for OCaml\")\n (description \"Composable, beautiful-by-default plotting built on Nx.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n (dune-configurator :build)\n (conf-sdl2 :build)\n (conf-cairo :build)\n (nx\n (= :version))\n (windtrap :with-test)\n (mdx :with-test))\n (tags\n (visualization plotting charts data-science graphics)))\n\n(package\n (name quill)\n (dir packages/quill)\n (synopsis \"Interactive REPL and markdown notebooks\")\n (description\n \"Quill is a REPL and notebook environment for OCaml. Interactive toplevel with syntax highlighting, completion, and history. Markdown notebooks with a terminal UI, web frontend, batch evaluation, and watch mode.\")\n (depends\n (ocaml\n (>= 5.2.0))\n dune\n cmarkit\n cmdliner\n bytesrw\n jsont\n mosaic\n (windtrap :with-test)\n (mdx :with-test))\n (tags\n (repl toplevel notebooks interactive-computing literate-programming)))\n\n(package\n (name raven)\n (allow_empty)\n (dir packages/raven)\n (synopsis \"Modern scientific computing for OCaml\")\n (description\n \"Raven is an ecosystem of composable libraries for numerical computing in OCaml. Tensors, automatic differentiation, neural networks, dataframes, plotting, tokenization, computer vision, reinforcement learning, and interactive notebooks.\")\n (depends\n (nx\n (= :version))\n (tolk\n (= :version))\n (brot\n (= :version))\n (talon\n (= :version))\n (rune\n (= :version))\n (vega\n (= :version))\n (kaun\n (= :version))\n (munin\n (= :version))\n (sowilo\n (= :version))\n (fehu\n (= :version))\n (hugin\n (= :version))\n (quill\n (= :version)))\n (tags\n (machine-learning data-science numerical-computation)))\n"
},
{
"path": "dune-workspace.tsan",
"content": "(lang dune 3.21)\n\n(lock_dir\n (path dune.lock))\n\n; Pin ocaml-variants to the 5.4 branch which includes the\n; __tsan_func_exit signature fix (ocaml/ocaml#14082).\n; Remove this pin once OCaml 5.4.2 is released.\n(pin\n (name ocaml-variants)\n (url \"git+https://github.com/ocaml/ocaml#5.4\")\n (package\n (name ocaml-variants)\n (version 5.4.2+trunk)))\n\n(lock_dir\n (path dune-tsan.lock)\n (pins ocaml-variants)\n (depopts ocaml-option-tsan))\n\n(context default)\n\n(context\n (default\n (name tsan)\n (lock_dir dune-tsan.lock)))\n"
},
{
"path": "opam/brot.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Tokenization for OCaml\"\ndescription:\n \"Fast, HuggingFace-compatible tokenization for language models. BPE, WordPiece, Unigram, word-level, and character-level algorithms with composable pipelines and training from scratch.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\n \"tokenization\" \"bpe\" \"wordpiece\" \"subword-tokenization\" \"language-models\"\n]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"re\"\n \"jsont\"\n \"bytesrw\"\n \"uunf\" {>= \"15.1.0\"}\n \"uucp\"\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/fehu.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Reinforcement learning for OCaml\"\ndescription:\n \"Type-safe RL environments, observation/action spaces, vectorized environments, trajectory collection, replay buffers, and generalized advantage estimation. Built on Nx.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"reinforcement-learning\" \"machine-learning\" \"environments\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"nx\" {= version}\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/hugin.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Declarative plotting and visualization for OCaml\"\ndescription: \"Composable, beautiful-by-default plotting built on Nx.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"visualization\" \"plotting\" \"charts\" \"data-science\" \"graphics\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"dune-configurator\" {build}\n \"conf-sdl2\" {build}\n \"conf-cairo\" {build}\n \"nx\" {= version}\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/kaun-board.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Training dashboard and logging for Raven\"\ndescription:\n \"Lightweight training logger and terminal dashboard for monitoring runs. The core library (kaun-board) provides a Log API for writing JSONL events and a reader for consuming them. The TUI library (kaun-board.tui) provides a Mosaic-based dashboard.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"training-dashboard\" \"monitoring\" \"logging\" \"machine-learning\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"jsont\"\n \"bytesrw\"\n \"cmdliner\"\n \"mosaic\"\n \"dune-configurator\" {build}\n \"matrix\"\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/kaun.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Neural networks for OCaml\"\ndescription:\n \"Composable layers, parameter trees, optimizers, training loops, data pipelines, and HuggingFace Hub integration. Built on Rune.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"neural-networks\" \"machine-learning\" \"deep-learning\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"rune\" {= version}\n \"vega\" {= version}\n \"nx\" {= version}\n \"jsont\"\n \"bytesrw\"\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/munin.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Local experiment tracking for Raven\"\ndescription:\n \"Local-first experiment tracking with append-only event logs, versioned artifacts, a terminal dashboard, and a CLI. The core library (munin) provides Session, Run, Store, and Artifact modules. The TUI library (munin.tui) provides a Mosaic-based dashboard.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"experiment-tracking\" \"machine-learning\" \"monitoring\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"jsont\"\n \"bytesrw\"\n \"sha\"\n \"cmdliner\"\n \"mosaic\"\n \"dune-configurator\" {build}\n \"matrix\"\n \"windtrap\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/norn.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"MCMC sampling for OCaml\"\ndescription:\n \"Markov chain Monte Carlo samplers with automatic gradients via Rune. Hamiltonian Monte Carlo with dual-averaging step-size adaptation.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"mcmc\" \"sampling\" \"bayesian\" \"machine-learning\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"nx\" {= version}\n \"rune\" {= version}\n \"windtrap\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/nx.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"N-dimensional arrays for OCaml\"\ndescription:\n \"Nx provides n-dimensional arrays with NumPy-like semantics and OCaml's type safety. 19 data types, broadcasting, slicing, linear algebra, FFT, and I/O. The numerical foundation for the Raven ecosystem.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"numerical-computation\" \"tensor-library\" \"machine-learning\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"dune-configurator\" {build}\n \"conf-pkg-config\" {build}\n \"conf-zlib\" {build}\n \"logs\"\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\ndepexts: [\n [\"libc-dev\" \"openblas-dev\" \"lapack-dev\"] {os-distribution = \"alpine\"}\n [\"epel-release\" \"openblas-devel\"] {os-distribution = \"centos\"}\n [\"libopenblas-dev\" \"liblapacke-dev\"] {os-family = \"debian\"}\n [\"libopenblas-dev\" \"liblapacke-dev\"] {os-family = \"ubuntu\"}\n [\"openblas-devel\"] {os-family = \"fedora\"}\n [\"libopenblas_openmp-devel\"] {os-family = \"suse\" | os-family = \"opensuse\"}\n [\"openblas\" \"lapacke\" \"cblas\"] {os-distribution = \"arch\"}\n [\"openblas\"] {os = \"macos\" & os-distribution = \"homebrew\"}\n [\"openblas\" \"lapacke\"] {os = \"freebsd\"}\n [\"mingw64-x86_64-cblas\" \"mingw64-x86_64-lapack\"] {os = \"cygwin\"}\n]\nx-ci-accept-failures: [\n \"oraclelinux-7\"\n \"oraclelinux-8\"\n \"oraclelinux-9\"\n]\n"
},
{
"path": "opam/nx.opam.template",
"content": "depexts: [\n [\"libc-dev\" \"openblas-dev\" \"lapack-dev\"] {os-distribution = \"alpine\"}\n [\"epel-release\" \"openblas-devel\"] {os-distribution = \"centos\"}\n [\"libopenblas-dev\" \"liblapacke-dev\"] {os-family = \"debian\"}\n [\"libopenblas-dev\" \"liblapacke-dev\"] {os-family = \"ubuntu\"}\n [\"openblas-devel\"] {os-family = \"fedora\"}\n [\"libopenblas_openmp-devel\"] {os-family = \"suse\" | os-family = \"opensuse\"}\n [\"openblas\" \"lapacke\" \"cblas\"] {os-distribution = \"arch\"}\n [\"openblas\"] {os = \"macos\" & os-distribution = \"homebrew\"}\n [\"openblas\" \"lapacke\"] {os = \"freebsd\"}\n [\"mingw64-x86_64-cblas\" \"mingw64-x86_64-lapack\"] {os = \"cygwin\"}\n]\nx-ci-accept-failures: [\n \"oraclelinux-7\"\n \"oraclelinux-8\"\n \"oraclelinux-9\"\n]\n"
},
{
"path": "opam/quill.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Interactive REPL and markdown notebooks\"\ndescription:\n \"Quill is a REPL and notebook environment for OCaml. Interactive toplevel with syntax highlighting, completion, and history. Markdown notebooks with a terminal UI, web frontend, batch evaluation, and watch mode.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\n \"repl\"\n \"toplevel\"\n \"notebooks\"\n \"interactive-computing\"\n \"literate-programming\"\n]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"cmarkit\"\n \"cmdliner\"\n \"bytesrw\"\n \"jsont\"\n \"mosaic\"\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/raven.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Modern scientific computing for OCaml\"\ndescription:\n \"Raven is an ecosystem of composable libraries for numerical computing in OCaml. Tensors, automatic differentiation, neural networks, dataframes, plotting, tokenization, computer vision, reinforcement learning, and interactive notebooks.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"machine-learning\" \"data-science\" \"numerical-computation\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"dune\" {>= \"3.21\"}\n \"nx\" {= version}\n \"tolk\" {= version}\n \"brot\" {= version}\n \"talon\" {= version}\n \"rune\" {= version}\n \"vega\" {= version}\n \"kaun\" {= version}\n \"munin\" {= version}\n \"sowilo\" {= version}\n \"fehu\" {= version}\n \"hugin\" {= version}\n \"quill\" {= version}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/rune.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Functional transformations for Nx arrays\"\ndescription:\n \"Automatic differentiation and vectorizing maps for Nx tensors. Reverse-mode AD (grad, vjp), forward-mode AD (jvp), vmap, and gradient checking, built on OCaml 5 effect handlers.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\n \"automatic-differentiation\"\n \"machine-learning\"\n \"deep-learning\"\n \"optimization\"\n]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"dune-configurator\" {build}\n \"nx\" {= version}\n \"tolk\" {= version}\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/sowilo.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Differentiable computer vision for OCaml\"\ndescription:\n \"Image processing operations expressed as Nx tensor computations. Geometric transforms, spatial filters, edge detection, morphological operations, and color space conversions, all compatible with Rune.grad and Rune.vmap.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\n \"computer-vision\" \"image-processing\" \"feature-detection\" \"machine-learning\"\n]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"nx\" {= version}\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/talon.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Dataframes for OCaml\"\ndescription:\n \"Fast and elegant dataframes with type-safe operations. Heterogeneous columns, applicative row operations, vectorized aggregations, and CSV I/O, built on Nx.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"dataframe\" \"data-manipulation\" \"data-science\" \"tabular-data\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"nx\" {= version}\n \"windtrap\" {with-test}\n \"mdx\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/tolk.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"A minimal ML compiler for GPU tensor computation\"\ndescription:\n \"Tolk is a minimal, readable ML compiler for GPU tensor computation in the Raven ecosystem.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"compiler\" \"gpu\" \"tensor-computation\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2\"}\n \"dune\" {>= \"3.21\"}\n \"windtrap\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "opam/vega.opam",
"content": "# This file is generated by dune, edit dune-project instead\nopam-version: \"2.0\"\nversion: \"1.0.0~alpha3\"\nsynopsis: \"Per-parameter gradient-based optimizers for OCaml\"\ndescription:\n \"Typed, per-parameter optimizer primitives: Adam, AdamW, SGD, RMSprop, Adagrad, and learning-rate schedules. Built on Nx with no autodiff dependency.\"\nmaintainer: [\"Thibaut Mattio \"]\nauthors: [\"Thibaut Mattio \"]\nlicense: \"ISC\"\ntags: [\"optimization\" \"machine-learning\" \"gradient-descent\"]\nhomepage: \"https://github.com/raven-ml/raven\"\ndoc: \"https://raven-ml.dev/docs/\"\nbug-reports: \"https://github.com/raven-ml/raven/issues\"\ndepends: [\n \"ocaml\" {>= \"5.2.0\"}\n \"dune\" {>= \"3.21\"}\n \"nx\" {= version}\n \"windtrap\" {with-test}\n \"thumper\" {with-test}\n \"odoc\" {with-doc}\n]\nbuild: [\n [\"dune\" \"subst\"] {dev}\n [\n \"dune\"\n \"build\"\n \"-p\"\n name\n \"-j\"\n jobs\n \"@install\"\n \"@runtest\" {with-test}\n \"@doc\" {with-doc}\n ]\n]\ndev-repo: \"git+https://github.com/raven-ml/raven.git\"\nx-maintenance-intent: [\"(latest)\"]\n"
},
{
"path": "packages/brot/README.md",
"content": "# Brot\n\nFast tokenization library for OCaml.\n\nBrot tokenizes text into token IDs for language models and reverses the\nprocess. It is part of the Raven ecosystem. It loads and saves HuggingFace\n`tokenizer.json` files, supports BPE, WordPiece, Unigram, word-level, and\ncharacter-level algorithms, and is 1.3-6x faster than HuggingFace\ntokenizers on most benchmarks.\n\n## Features\n\n- Tokenization algorithms: BPE, WordPiece, Unigram, word-level, character-level\n- HuggingFace compatible: load and save `tokenizer.json` files, load\n vocab/merges model files\n- Composable pipeline: normalizer, pre-tokenizer, post-processor, decoder\n — each stage independently configurable\n- Rich encoding output: token IDs, string tokens, byte offsets, attention\n masks, type IDs, word IDs, special token masks\n- Training: train BPE, WordPiece, Unigram, and word-level tokenizers from\n scratch\n- Performance: 1.3-6x faster than HuggingFace tokenizers (Rust native) on\n most benchmarks — see [bench/](bench/) for details\n\n## Quick Start\n\n\n```ocaml\nopen Brot\n\nlet () =\n (* Load a pretrained HuggingFace tokenizer *)\n let tokenizer = from_file \"tokenizer.json\" |> Result.get_ok in\n\n (* Encode text to token IDs *)\n let encoding = encode tokenizer \"Hello world!\" in\n let ids = Encoding.ids encoding in\n Printf.printf \"Token IDs: \";\n Array.iter (fun id -> Printf.printf \"%d \" id) ids;\n print_newline ();\n\n (* Decode back to text *)\n let text = decode tokenizer ids in\n Printf.printf \"Decoded: %s\\n\" text\n```\n\n## Contributing\n\nSee the [Raven monorepo README](../README.md) for contribution guidelines.\n\n## License\n\nISC License. See [LICENSE](../LICENSE) for details.\n"
},
{
"path": "packages/brot/bench/README.md",
"content": "# Brot Benchmarks\n\nThis directory contains micro-benchmarks for the `brot` library.\nThe suite mirrors HuggingFace's `tokenizers` so we can compare wall-clock\nthroughput for realistic workloads and catch regressions.\n\n## Fixtures\n\nBenchmark inputs live in `./data/`:\n\n- `news_1k.txt`, `wiki_64k.txt`, `code_excerpt.txt` — sample corpora used for\n encoding workloads.\n- `gpt2.json` — OpenAI GPT-2 (BPE, 50K vocab, 50K merges)\n- `bert_base.json` — Google BERT-base-uncased (WordPiece, 30K vocab)\n- `llama.json` — Meta LLaMA (BPE, 32K vocab, 61K merges, no pre-tokenizer)\n\nDownload the tokenizer model files:\n\n```bash\nbrot/bench/download_data.sh\n```\n\n## Running the Benchmarks\n\n### Brot (OCaml)\n\n```bash\ndune exec brot/bench/bench_brot.exe -- --gc\n```\n\n### tokenizers — Rust native\n\n```bash\ncd brot/bench/bench_rust && cargo run --release\n```\n\n### tokenizers — Python (Rust FFI)\n\n```bash\nuv run --with tokenizers brot/bench/bench_tokenizers.py\n```\n\n## Comparison\n\nWall-clock time per run. Lower is better. Apple M3 Pro, macOS.\n\n| Benchmark | Brot (OCaml) | Rust native | Python (Rust FFI) | Brot vs Rust |\n| ------------------------------------ | ------------ | ----------- | ----------------- | ------------ |\n| **GPT-2** (BPE, 50K vocab) | | | | |\n| Encode/short (1KB) | 46μs | 209μs | 250μs | **4.5x** |\n| Encode/long (64KB) | 5.26ms | 10.25ms | 13.27ms | **1.9x** |\n| Encode/batch_32 | 1.38ms | 3.05ms | 3.91ms | **2.2x** |\n| Decode/long | 1.19ms | 1.50ms | 1.58ms | **1.3x** |\n| **BERT-base** (WordPiece, 30K vocab) | | | | |\n| Encode/short (1KB) | 137μs | 278μs | 325μs | **2.0x** |\n| Encode/long (64KB) | 10.87ms | 13.95ms | 16.64ms | **1.3x** |\n| Encode/batch_32 | 2.06ms | 2.31ms | 2.66ms | **1.1x** |\n| Decode/long | 1.25ms | 7.63ms | 7.76ms | **6.1x** |\n| **LLaMA** (BPE, 32K vocab) | | | | |\n| Encode/short (1KB) | 51μs | 207μs | 247μs | **4.1x** |\n| Encode/long (64KB) | 20.15ms | 13.41ms | 16.23ms | 1.5x slower |\n| Encode/batch_32 | 1.43ms | 1.56ms | 1.51ms | ~par |\n| Decode/long | 1.12ms | 5.02ms | 5.03ms | **4.5x** |\n\nNotes:\n- The \"Rust native\" column calls the `tokenizers` crate directly, no Python FFI.\n Source: `bench_rust/main.rs`.\n- Both brot and HF tokenizers use multi-threading for batch encoding (wall < CPU).\n- LLaMA has no pre-tokenizer, so the entire text goes through BPE as a single\n sequence — this is where brot's BPE is slower on long inputs.\n"
},
{
"path": "packages/brot/bench/bench_brot.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Benchmark suite for Brot tokenizers using realistic fixtures. *)\n\nopen Brot\n\nmodule Fixtures = struct\n let data_dir = Filename.concat (Sys.getcwd ()) \"packages/brot/bench/data\"\n\n let read_file name =\n let path = Filename.concat data_dir name in\n let ic = open_in_bin path in\n let len = in_channel_length ic in\n let content = really_input_string ic len in\n close_in ic;\n content\n\n let load_tokenizer name =\n let path = Filename.concat data_dir name in\n match from_file path with\n | Ok tok -> tok\n | Error msg ->\n failwith (Printf.sprintf \"Failed to load tokenizer %s: %s\" path msg)\n\n let short_text = read_file \"news_1k.txt\"\n let long_text = read_file \"wiki_64k.txt\"\n\n let batch_32 =\n let rec loop acc remaining =\n if remaining = 0 then List.rev acc\n else loop (short_text :: acc) (remaining - 1)\n in\n loop [] 32\nend\n\nlet encode_single tok text = encode tok text\nlet encode_batch tok texts = encode_batch tok texts\nlet decode_ids tok ids = decode tok ids\n\nlet make_suite ~label ~tokenizer =\n let open Fixtures in\n let decode_input =\n let encoding = encode tokenizer long_text in\n Array.copy (Encoding.ids encoding)\n in\n let benches =\n [\n Thumper.bench \"Encode/single_short\" (fun () ->\n encode_single tokenizer short_text);\n Thumper.bench \"Encode/single_long\" (fun () ->\n encode_single tokenizer long_text);\n Thumper.bench \"Encode/batch_32\" (fun () ->\n encode_batch tokenizer batch_32);\n Thumper.bench \"Decode/long\" (fun () -> decode_ids tokenizer decode_input);\n ]\n in\n Thumper.group label benches\n\nlet all_benchmarks =\n let open Fixtures in\n let gpt2 =\n make_suite ~label:\"GPT-2\" ~tokenizer:(load_tokenizer \"gpt2.json\")\n in\n let bert =\n make_suite ~label:\"BERT-base\" ~tokenizer:(load_tokenizer \"bert_base.json\")\n in\n let llama =\n make_suite ~label:\"LLaMA\" ~tokenizer:(load_tokenizer \"llama.json\")\n in\n [ gpt2; bert; llama ]\n\nlet () = Thumper.run \"brot\" all_benchmarks\n"
},
{
"path": "packages/brot/bench/bench_rust/.gitignore",
"content": "/target\n"
},
{
"path": "packages/brot/bench/bench_rust/Cargo.toml",
"content": "[package]\nname = \"bench_tokenizers_rust\"\nedition = \"2021\"\n\n[[bin]]\nname = \"bench_tokenizers_rust\"\npath = \"main.rs\"\n\n[dependencies]\ntokenizers = \"0.22\"\n"
},
{
"path": "packages/brot/bench/bench_rust/main.rs",
"content": "use std::fs;\nuse std::path::Path;\nuse std::time::{Duration, Instant};\nuse tokenizers::Tokenizer;\n\nconst WARMUP: usize = 4;\nconst TIME_QUOTA: Duration = Duration::from_millis(300);\nconst MIN_MEASUREMENTS: usize = 3;\n\nstruct BenchResult {\n name: String,\n wall_per_run: Duration,\n runs: usize,\n}\n\nfn bench(name: &str, mut f: F) -> BenchResult {\n // Warmup\n for _ in 0..WARMUP {\n f();\n }\n\n // Adaptive batching: start with batch_size=1, scale up until each batch\n // takes at least 2ms of wall time, then collect measurements for ~0.3s.\n let mut batch_size: usize = 1;\n let mut measurements: Vec = Vec::new();\n let bench_start = Instant::now();\n\n loop {\n let start = Instant::now();\n for _ in 0..batch_size {\n f();\n }\n let elapsed = start.elapsed();\n\n if elapsed.as_secs_f64() < 0.002 {\n // Batch too fast, scale up\n batch_size = (batch_size as f64 * 1.3).ceil().max((batch_size + 1) as f64) as usize;\n continue;\n }\n\n let per_run = elapsed / batch_size as u32;\n measurements.push(per_run);\n\n let total_elapsed = bench_start.elapsed();\n if measurements.len() >= MIN_MEASUREMENTS && total_elapsed >= TIME_QUOTA {\n break;\n }\n\n batch_size = (batch_size as f64 * 1.3).ceil().max((batch_size + 1) as f64) as usize;\n }\n\n // Compute average\n let total: Duration = measurements.iter().sum();\n let avg = total / measurements.len() as u32;\n\n BenchResult {\n name: name.to_string(),\n wall_per_run: avg,\n runs: measurements.len(),\n }\n}\n\nfn format_duration(d: Duration) -> String {\n let nanos = d.as_nanos() as f64;\n if nanos < 1_000.0 {\n format!(\"{:.2}ns\", nanos)\n } else if nanos < 1_000_000.0 {\n format!(\"{:.2}μs\", nanos / 1_000.0)\n } else if nanos < 1_000_000_000.0 {\n format!(\"{:.2}ms\", nanos / 1_000_000.0)\n } else {\n format!(\"{:.2}s\", nanos / 1_000_000_000.0)\n }\n}\n\nfn run_suite(label: &str, tokenizer: &Tokenizer, short_text: &str, long_text: &str) {\n let batch_32: Vec<&str> = vec![short_text; 32];\n\n // Pre-compute decode input\n let encoding = tokenizer\n .encode(long_text, false)\n .expect(\"encode for decode input\");\n let decode_ids: Vec = encoding.get_ids().to_vec();\n\n let results = vec![\n bench(&format!(\"{}/Encode/single_short\", label), || {\n tokenizer.encode(short_text, false).unwrap();\n }),\n bench(&format!(\"{}/Encode/single_long\", label), || {\n tokenizer.encode(long_text, false).unwrap();\n }),\n bench(&format!(\"{}/Encode/batch_32\", label), || {\n tokenizer\n .encode_batch(batch_32.clone(), false)\n .unwrap();\n }),\n bench(&format!(\"{}/Decode/long\", label), || {\n tokenizer.decode(decode_ids.as_slice(), true).unwrap();\n }),\n ];\n\n for r in &results {\n println!(\n \" {:<35} {:>10} ({} samples)\",\n r.name,\n format_duration(r.wall_per_run),\n r.runs\n );\n }\n}\n\nfn main() {\n let data_dir = Path::new(env!(\"CARGO_MANIFEST_DIR\")).join(\"../data\");\n\n let short_text =\n fs::read_to_string(data_dir.join(\"news_1k.txt\")).expect(\"read news_1k.txt\");\n let long_text =\n fs::read_to_string(data_dir.join(\"wiki_64k.txt\")).expect(\"read wiki_64k.txt\");\n\n println!(\"Rust-native HuggingFace tokenizers benchmark\");\n println!(\"=============================================\\n\");\n\n let tokenizers = [\n (\"GPT-2\", \"gpt2.json\"),\n (\"BERT-base\", \"bert_base.json\"),\n (\"LLaMA\", \"llama.json\"),\n ];\n\n for (label, filename) in &tokenizers {\n let path = data_dir.join(filename);\n let tokenizer =\n Tokenizer::from_file(&path).unwrap_or_else(|e| {\n panic!(\"Failed to load {}: {}\", path.display(), e)\n });\n println!(\"{}:\", label);\n run_suite(label, &tokenizer, &short_text, &long_text);\n println!();\n }\n}\n"
},
{
"path": "packages/brot/bench/bench_tokenizers.py",
"content": "from __future__ import annotations\n\nfrom pathlib import Path\nfrom typing import Any, Callable, List\n\nfrom tokenizers import Tokenizer\n\n_ROOT = Path(__file__).resolve().parent\n_DATA_DIR = _ROOT / \"data\"\n\nimport sys\n\n_SCRIPTS_DIR = _ROOT\nwhile not (_SCRIPTS_DIR / \"dune-project\").exists():\n _SCRIPTS_DIR = _SCRIPTS_DIR.parent\n_SCRIPTS_DIR = _SCRIPTS_DIR / \"scripts\"\nif str(_SCRIPTS_DIR) not in sys.path:\n sys.path.insert(0, str(_SCRIPTS_DIR))\n\nimport ubench # type: ignore\n\n\nSHORT_TEXT = (_DATA_DIR / \"news_1k.txt\").read_text(encoding=\"utf-8\")\nLONG_TEXT = (_DATA_DIR / \"wiki_64k.txt\").read_text(encoding=\"utf-8\")\nBATCH_32 = [SHORT_TEXT] * 32\n\n\ndef load_tokenizer(filename: str) -> Tokenizer:\n path = _DATA_DIR / filename\n return Tokenizer.from_file(str(path))\n\n\ndef make_suite(label: str, tokenizer: Tokenizer) -> Any:\n decode_ids = tokenizer.encode(LONG_TEXT).ids\n\n benches: List[Any] = [\n ubench.bench(\"Encode/single_short\", lambda: tokenizer.encode(SHORT_TEXT)),\n ubench.bench(\"Encode/single_long\", lambda: tokenizer.encode(LONG_TEXT)),\n ubench.bench(\"Encode/batch_32\", lambda: tokenizer.encode_batch(BATCH_32)),\n ubench.bench(\"Decode/long\", lambda: tokenizer.decode(decode_ids)),\n ]\n\n return ubench.group(label, benches)\n\n\ndef build_benchmarks() -> List[Any]:\n return [\n make_suite(\"GPT-2\", load_tokenizer(\"gpt2.json\")),\n make_suite(\"BERT-base\", load_tokenizer(\"bert_base.json\")),\n make_suite(\"LLaMA\", load_tokenizer(\"llama.json\")),\n ]\n\n\ndef default_config() -> ubench.Config:\n return ubench.Config.default().build()\n\n\ndef main() -> None:\n benchmarks = build_benchmarks()\n config = default_config()\n ubench.run(benchmarks, config=config, output_format=\"pretty\", verbose=False)\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "packages/brot/bench/brot.thumper",
"content": "# thumper baseline\n# version: 1\n# suite_name: brot\n# host: 1480401c3b76ed18\n# cpu: Apple M1 Max\n# ocaml: 5.4.1\n# git: 31747323\n# dirty: true\n# command: /Users/tmattio/Workspace/raven/_build/default/packages/brot/bench/bench_brot.exe --bless --quick\n\nbert-base/decode_long\talloc_words\t4.445400e+05\t4.445400e+05\t4.445400e+05\t0.000000e+00\t9\t0\nbert-base/decode_long\tcpu_time\t1.424685e-03\t1.370913e-03\t1.470915e-03\t3.509645e-02\t9\t1\nbert-base/decode_long\twall_time\t1.425388e-03\t1.371378e-03\t1.475911e-03\t3.666840e-02\t9\t1\nbert-base/encode_batch_32\talloc_words\t1.089250e+05\t1.089250e+05\t1.089250e+05\t0.000000e+00\t31\t1\nbert-base/encode_batch_32\tcpu_time\t9.091263e-03\t8.686646e-03\t9.591658e-03\t4.977371e-02\t31\t2\nbert-base/encode_batch_32\twall_time\t2.121498e-03\t1.993592e-03\t2.253395e-03\t6.123112e-02\t31\t1\nbert-base/encode_single_long\talloc_words\t1.350547e+06\t1.350547e+06\t1.350547e+06\t0.000000e+00\t9\t1\nbert-base/encode_single_long\tcpu_time\t9.506189e-03\t9.354674e-03\t9.692541e-03\t1.777085e-02\t9\t1\nbert-base/encode_single_long\twall_time\t9.509793e-03\t9.372449e-03\t9.680291e-03\t1.618554e-02\t9\t1\nbert-base/encode_single_short\talloc_words\t2.699600e+04\t2.699600e+04\t2.699600e+04\t0.000000e+00\t9\t0\nbert-base/encode_single_short\tcpu_time\t1.392726e-04\t1.345864e-04\t1.448241e-04\t3.675399e-02\t9\t0\nbert-base/encode_single_short\twall_time\t1.393180e-04\t1.345034e-04\t1.440607e-04\t3.430033e-02\t9\t0\ngpt-2/decode_long\talloc_words\t3.417770e+05\t3.417770e+05\t3.417770e+05\t0.000000e+00\t7\t0\ngpt-2/decode_long\tcpu_time\t1.305595e-03\t1.261472e-03\t1.338931e-03\t2.966443e-02\t7\t0\ngpt-2/decode_long\twall_time\t1.305703e-03\t1.262279e-03\t1.346462e-03\t3.223653e-02\t7\t0\ngpt-2/encode_batch_32\talloc_words\t5.518900e+04\t5.518900e+04\t5.518900e+04\t0.000000e+00\t6\t0\ngpt-2/encode_batch_32\tcpu_time\t3.952923e-03\t3.848310e-03\t4.113309e-03\t3.351934e-02\t6\t1\ngpt-2/encode_batch_32\twall_time\t1.386324e-03\t1.328412e-03\t1.431061e-03\t3.702195e-02\t6\t0\ngpt-2/encode_single_long\talloc_words\t6.731690e+05\t6.731690e+05\t6.731690e+05\t0.000000e+00\t19\t2\ngpt-2/encode_single_long\tcpu_time\t3.852835e-03\t3.758677e-03\t3.927922e-03\t2.196376e-02\t19\t1\ngpt-2/encode_single_long\twall_time\t3.856248e-03\t3.756035e-03\t3.930090e-03\t2.256793e-02\t19\t1\ngpt-2/encode_single_short\talloc_words\t1.356200e+04\t1.356200e+04\t1.356200e+04\t0.000000e+00\t13\t0\ngpt-2/encode_single_short\tcpu_time\t5.279107e-05\t5.090596e-05\t5.553182e-05\t4.381283e-02\t13\t0\ngpt-2/encode_single_short\twall_time\t5.282309e-05\t5.106692e-05\t5.499492e-05\t3.718073e-02\t13\t0\nllama/decode_long\talloc_words\t6.844460e+05\t6.844460e+05\t6.844460e+05\t0.000000e+00\t12\t0\nllama/decode_long\tcpu_time\t1.149214e-03\t1.094901e-03\t1.194180e-03\t4.319437e-02\t12\t0\nllama/decode_long\twall_time\t1.149682e-03\t1.095042e-03\t1.189211e-03\t4.095449e-02\t12\t0\nllama/encode_batch_32\talloc_words\t9.471700e+04\t9.471700e+04\t9.471700e+04\t0.000000e+00\t5\t0\nllama/encode_batch_32\tcpu_time\t4.421498e-03\t4.320283e-03\t4.534447e-03\t2.421853e-02\t5\t0\nllama/encode_batch_32\twall_time\t1.467702e-03\t1.366193e-03\t1.593832e-03\t7.754944e-02\t5\t2\nllama/encode_single_long\talloc_words\t1.261210e+06\t1.261210e+06\t1.261210e+06\t0.000000e+00\t9\t1\nllama/encode_single_long\tcpu_time\t1.817278e-02\t1.788729e-02\t1.860684e-02\t1.979736e-02\t9\t2\nllama/encode_single_long\twall_time\t1.819150e-02\t1.794107e-02\t1.863329e-02\t1.902594e-02\t9\t2\nllama/encode_single_short\talloc_words\t2.344400e+04\t2.344400e+04\t2.344400e+04\t0.000000e+00\t42\t0\nllama/encode_single_short\tcpu_time\t6.126139e-05\t6.069695e-05\t6.174501e-05\t8.553960e-03\t42\t6\nllama/encode_single_short\twall_time\t6.130598e-05\t6.079485e-05\t6.183071e-05\t8.448240e-03\t42\t6\n"
},
{
"path": "packages/brot/bench/data/.gitignore",
"content": "gpt2.json\nbert_base.json\nllama.json\n"
},
{
"path": "packages/brot/bench/data/news_1k.txt",
"content": "City officials confirmed on Tuesday that the riverside park will reopen this summer after a two-year renovation.\nCrews installed 175 energy-efficient lights, replanted native wildflowers, and added a playground designed by local artists.\nThe project ran $1.8 million under budget, according to Deputy Mayor Alicia Gómez — a welcome surprise for residents concerned about rising taxes.\n\"It's not just a facelift; it's a commitment to public space,\" said Gomez. Cyclists tested the new bike lanes, while children chased bubbles during the ribbon-cutting ceremony.\nThe park will host weekly night markets featuring Afghan bolani, Jamaican patties, and vegan empanadas, with vendors selected through a community ballot.\nPublic transit advocates noted that the expanded bus schedule, combined with real-time arrival boards, should alleviate weekend congestion.\nSustainability officers also unveiled a solar-powered irrigation system and a pollinator habitat that includes milkweed, lavender, and rare prairie clover.\nEarly visitor surveys show 92% satisfaction, with many praising the accessible design, tactile maps, and multilingual audio tours available in English, Spanish, Mandarin, and American Sign Language.\nThe city plans to share open-source blueprints and a detailed maintenance playbook with other municipalities considering similar upgrades."
},
{
"path": "packages/brot/bench/data/wiki_64k.txt",
"content": "== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday Life ==\nSchoolchildren log phenology observations, while retired tram conductors teach visitor orientation classes in a repurposed depot, complete with time-travel escape room puzzles chronicling the town's evolution.\n\n== Early History ==\n\nThe settlement traces its roots to a trading village documented in the 12th-century annals of the Seljuk chronicler al-Biruni. Archaeological digs in 1989 uncovered kiln-fired ceramics, copper ingots, and terraced irrigation canals that reshaped historians' understanding of Central Asian trade routes.\n\n== Linguistics ==\nModern dialect surveys reveal a blend of Chuvash, Khazar, and Oghur loanwords; linguists have mapped palatalized consonants appearing near river valleys, likely a relic of seasonal migration.\n== Technological Renaissance ==\nBy 1893 the town hosted one of the earliest wireless telegraph stations in the region. Engineer Lidiya Petrovna retrofitted surplus naval equipment to send meteorological data to Moscow every sunset. Her notebooks — digitized in 2017 — contain meticulous diagrams of spark-gap transmitters, annotations in French, and the occasional doodle of a cat wearing goggles.\n== Cultural Revival ==\nAnnual festivals now feature Tuvan throat singing workshops, VR reconstructions of vanished monasteries, and fermentation labs explaining the chemistry behind kumis. UNESCO added the town's accordion workshops to its intangible heritage list, citing their adaptive use of recycled polymers for reeds.\n== Contemporary Research ==\nIn 2022 a consortium of botanists, data journalists, and Indigenous seed keepers launched the Steppe Observatory, using open satellite data, LoRaWAN sensors, and community weather diaries to forecast dust storms.\n== Notable Figures ==\nHistorian Salome Okafor popularized the settlement after translating 400 folktales into Yoruba, English, and Esperanto, each annotated with QR codes linking to oral history recordings.\n== Gastronomy ==\nLocal chefs pair fermented camel-milk cheese with candied sea buckthorn, while food trucks experiment with kelp-laden naan tacos, reflecting the town's fishing diaspora.\n== Climate Adaptation ==\nFlood mitigation now involves mycelium-reinforced levees, willow microforests, and AI-optimized sluice gates governed by a civic algorithm crafted in nightly town halls.\n== Digital Archives ==\nVolunteer coders maintain a mirrored archive stored on solar-powered Raspberry Pi clusters. The archive syncs monthly via a community-owned satellite uplink leased during lunar downtimes.\n== Everyday L"
},
{
"path": "packages/brot/bench/download_data.sh",
"content": "#!/usr/bin/env bash\nset -euo pipefail\n\nDATA_DIR=\"$(cd \"$(dirname \"$0\")/data\" && pwd)\"\n\necho \"Downloading real-world tokenizer models to $DATA_DIR...\"\n\ncurl -sL -o \"$DATA_DIR/gpt2.json\" \\\n \"https://huggingface.co/openai-community/gpt2/resolve/main/tokenizer.json\"\necho \" GPT-2 (BPE, 50K vocab)\"\n\ncurl -sL -o \"$DATA_DIR/bert_base.json\" \\\n \"https://huggingface.co/google-bert/bert-base-uncased/resolve/main/tokenizer.json\"\necho \" BERT-base (WordPiece, 30K vocab)\"\n\ncurl -sL -o \"$DATA_DIR/llama.json\" \\\n \"https://huggingface.co/hf-internal-testing/llama-tokenizer/resolve/main/tokenizer.json\"\necho \" LLaMA (BPE, 32K vocab)\"\n\necho \"Done.\"\n"
},
{
"path": "packages/brot/bench/dune",
"content": "(data_only_dirs data)\n\n(executable\n (name bench_brot)\n (libraries brot thumper unix))\n\n(rule\n (alias runtest)\n (action\n (progn\n (run %{exe:bench_brot.exe} -q)\n (diff? brot.thumper brot.thumper.corrected))))\n"
},
{
"path": "packages/brot/doc/01-getting-started.md",
"content": "# Getting Started\n\nThis guide covers the basics: encoding text to token IDs, decoding back\nto text, configuring the pipeline, and training tokenizers from scratch.\n\n## Installation\n\n\n```bash\nopam install brot\n```\n\nOr build from source:\n\n\n```bash\ngit clone https://github.com/raven-ml/raven\ncd raven && dune build brot\n```\n\n## Encoding and Decoding\n\nA tokenizer converts text to token IDs and back. Build one from a\nvocabulary and merge rules, then encode and decode:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n bpe\n ~vocab:\n [ (\"h\", 0); (\"e\", 1); (\"l\", 2); (\"o\", 3); (\" \", 4); (\"w\", 5);\n (\"r\", 6); (\"d\", 7); (\"he\", 8); (\"ll\", 9); (\"llo\", 10);\n (\"hello\", 11); (\"wo\", 12); (\"rl\", 13); (\"rld\", 14); (\"world\", 15) ]\n ~merges:\n [ (\"h\", \"e\"); (\"l\", \"l\"); (\"ll\", \"o\"); (\"he\", \"llo\");\n (\"w\", \"o\"); (\"r\", \"l\"); (\"rl\", \"d\"); (\"wo\", \"rld\") ]\n ()\n\n(* Encode text to an Encoding *)\nlet encoding = encode tokenizer \"hello world\"\nlet ids = Encoding.ids encoding (* [| 11; 4; 15 |] *)\nlet tokens = Encoding.tokens encoding (* [| \"hello\"; \" \"; \"world\" |] *)\n\n(* Decode back to text *)\nlet text = decode tokenizer ids (* \"hello world\" *)\n```\n\n`encode` returns an `Encoding.t`. For just the IDs, use `encode_ids`:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n bpe\n ~vocab:\n [ (\"h\", 0); (\"e\", 1); (\"l\", 2); (\"o\", 3); (\" \", 4); (\"w\", 5);\n (\"r\", 6); (\"d\", 7); (\"he\", 8); (\"ll\", 9); (\"llo\", 10);\n (\"hello\", 11); (\"wo\", 12); (\"rl\", 13); (\"rld\", 14); (\"world\", 15) ]\n ~merges:\n [ (\"h\", \"e\"); (\"l\", \"l\"); (\"ll\", \"o\"); (\"he\", \"llo\");\n (\"w\", \"o\"); (\"r\", \"l\"); (\"rl\", \"d\"); (\"wo\", \"rld\") ]\n ()\n\nlet ids = encode_ids tokenizer \"hello world\" (* [| 11; 4; 15 |] *)\n```\n\n## Encoding Output\n\nAn `Encoding.t` carries more than just token IDs. Every field is a\nparallel array of the same length:\n\n- `ids` — integer token IDs for model input\n- `tokens` — string representation of each token\n- `offsets` — `(start, end)` byte positions in the original text\n- `type_ids` — segment IDs (0 for first sentence, 1 for second in pair tasks)\n- `attention_mask` — 1 for real tokens, 0 for padding\n- `special_tokens_mask` — 1 for special tokens (`[CLS]`, `[SEP]`, padding), 0 for content\n- `word_ids` — maps each token to its source word index, or `None` for special tokens\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[UNK]\", 0); (\"[CLS]\", 1); (\"[SEP]\", 2);\n (\"the\", 3); (\"cat\", 4); (\"play\", 5); (\"##ing\", 6) ]\n ~specials:(List.map special [ \"[UNK]\"; \"[CLS]\"; \"[SEP]\" ])\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) ())\n ~decoder:(Decoder.wordpiece ())\n ~pre:(Pre_tokenizer.whitespace ())\n ~unk_token:\"[UNK]\" ()\n\nlet enc = encode tokenizer \"the cat playing\"\n(* tokens: [| \"[CLS]\"; \"the\"; \"cat\"; \"play\"; \"##ing\"; \"[SEP]\" |] *)\nlet ids = Encoding.ids enc\nlet type_ids = Encoding.type_ids enc\nlet attention_mask = Encoding.attention_mask enc\nlet special_tokens_mask = Encoding.special_tokens_mask enc\nlet offsets = Encoding.offsets enc\nlet word_ids = Encoding.word_ids enc\n```\n\nSee [Batch Processing](04-batch-processing/) for a deeper look at encoding\nmetadata, sentence pairs, padding, and truncation.\n\n## The Pipeline\n\nTokenization proceeds through up to 5 configurable stages:\n\n1. **Normalizer** — text cleanup (lowercase, accent removal, Unicode normalization)\n2. **Pre-tokenizer** — split text into pieces with byte offsets\n3. **Algorithm** — apply vocabulary-based encoding (BPE, WordPiece, Unigram, etc.)\n4. **Post-processor** — add special tokens and set type IDs\n5. **Decoder** — reverse the encoding back to text\n\nEach stage is optional. Here is a complete BERT-style pipeline:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~normalizer:(Normalizer.bert ~lowercase:true ())\n ~pre:(Pre_tokenizer.bert ())\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) ())\n ~decoder:(Decoder.wordpiece ())\n ~vocab:\n [ (\"[UNK]\", 0); (\"[CLS]\", 1); (\"[SEP]\", 2); (\"[PAD]\", 3);\n (\"the\", 4); (\"cat\", 5); (\"sat\", 6); (\"on\", 7);\n (\"play\", 8); (\"##ing\", 9); (\"##ed\", 10) ]\n ~specials:(List.map special [ \"[UNK]\"; \"[CLS]\"; \"[SEP]\"; \"[PAD]\" ])\n ~unk_token:\"[UNK]\" ~pad_token:\"[PAD]\" ()\n\n(* The normalizer lowercases \"The Cat\" before tokenization *)\nlet enc = encode tokenizer \"The Cat Sat\"\nlet tokens = Encoding.tokens enc\n(* [| \"[CLS]\"; \"the\"; \"cat\"; \"sat\"; \"[SEP]\" |] *)\n\n(* Decode, skipping special tokens *)\nlet text = decode tokenizer ~skip_special_tokens:true (Encoding.ids enc)\n(* \"the cat sat\" *)\n```\n\nSee [The Tokenization Pipeline](02-pipeline/) for a detailed guide to each\nstage.\n\n## Training\n\nTrain a tokenizer from a text corpus. Brot supports training BPE,\nWordPiece, Unigram, and word-level tokenizers:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n train_bpe ~vocab_size:80 ~show_progress:false\n (`Seq (List.to_seq\n [ \"The quick brown fox jumps over the lazy dog\";\n \"The dog barked loudly at the brown fox\";\n \"Quick brown foxes are jumping over lazy dogs\";\n \"The lazy dog slept while the fox jumped\" ]))\n\nlet size = vocab_size tokenizer\nlet enc = encode tokenizer \"The quick fox\"\n```\n\nSee [Choosing an Algorithm](05-algorithms/) for guidance on which algorithm\nto use and how to configure training.\n\n## Loading Pretrained Tokenizers\n\nLoad a HuggingFace `tokenizer.json` file:\n\n\n```ocaml\nopen Brot\n\nlet tokenizer = from_file \"tokenizer.json\" |> Result.get_ok\nlet encoding = encode tokenizer \"Hello world!\"\n```\n\nLoad from separate vocabulary and merges files:\n\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n from_model_file ~vocab:\"vocab.json\" ~merges:\"merges.txt\"\n ~pre:(Pre_tokenizer.byte_level ~add_prefix_space:false ())\n ~decoder:(Decoder.byte_level ())\n ()\n```\n\nSee [Pretrained Tokenizers](03-pretrained/) for complete pipeline\nconfigurations for BERT, GPT-2, and SentencePiece-style models.\n\n## Batch Processing\n\nEncode multiple texts at once with padding to uniform length:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n train_bpe ~vocab_size:80 ~show_progress:false\n ~specials:(List.map special [ \"[PAD]\" ])\n ~pad_token:\"[PAD]\"\n (`Seq (List.to_seq\n [ \"The quick brown fox jumps over the lazy dog\";\n \"The dog barked loudly at the brown fox\";\n \"Quick brown foxes are jumping over lazy dogs\" ]))\n\nlet encodings =\n encode_batch tokenizer\n ~padding:(padding `Batch_longest)\n [ \"The quick fox\"; \"The lazy dog barked\" ]\n\n(* All encodings now have the same length *)\nlet lengths = List.map Encoding.length encodings\n```\n\nSee [Batch Processing](04-batch-processing/) for padding strategies,\ntruncation, sentence pairs, and offset alignment.\n\n## Next Steps\n\n- [The Tokenization Pipeline](02-pipeline/) — how the 5 pipeline stages work\n- [Pretrained Tokenizers](03-pretrained/) — loading, saving, and building known model pipelines\n- [Batch Processing](04-batch-processing/) — padding, truncation, encoding metadata\n- [Choosing an Algorithm](05-algorithms/) — BPE vs WordPiece vs Unigram and when to use each\n"
},
{
"path": "packages/brot/doc/02-pipeline.md",
"content": "# The Tokenization Pipeline\n\nBrot processes text through up to 5 stages, each optional and independently\nconfigurable:\n\n```\ntext\n │\n ├─ 1. Normalizer — clean and transform text\n ├─ 2. Pre-tokenizer — split into pieces with byte offsets\n ├─ 3. Algorithm — map pieces to token IDs (BPE, WordPiece, …)\n ├─ 4. Post-processor — add special tokens, set type IDs\n └─ 5. Decoder — reverse the encoding back to text\n │\n ▼\nEncoding.t (ids, tokens, offsets, masks, …)\n```\n\nEach stage is set when constructing the tokenizer. Omit any stage and it\nis skipped.\n\n## Normalization\n\nNormalizers transform text before tokenization. They handle lowercasing,\naccent removal, Unicode normalization, whitespace cleanup, and\nmodel-specific preprocessing.\n\nAvailable normalizers:\n\n- **Unicode**: `nfc`, `nfd`, `nfkc`, `nfkd`\n- **Text transforms**: `lowercase`, `strip_accents`, `strip`, `replace`, `prepend`\n- **Byte-level**: `byte_level` (GPT-2 style byte-to-Unicode mapping)\n- **Model-specific**: `bert` (clean text, CJK padding, optional lowercasing and accent stripping)\n\nCompose normalizers with `sequence`:\n\n```ocaml\nopen Brot\n\nlet n =\n Normalizer.sequence\n [ Normalizer.nfd; Normalizer.strip_accents; Normalizer.lowercase ]\n\nlet r1 = Normalizer.apply n \"Café Résumé\" (* \"cafe resume\" *)\nlet r2 = Normalizer.apply n \"HELLO\" (* \"hello\" *)\n```\n\nThe BERT normalizer combines several transforms:\n\n```ocaml\nopen Brot\n\nlet n = Normalizer.bert ~lowercase:true ()\n(* Lowercases, cleans control characters, pads CJK *)\nlet r1 = Normalizer.apply n \"Hello World\" (* \"hello world\" *)\nlet r2 = Normalizer.apply n \"Café\" (* \"cafe\" *)\n```\n\n## Pre-tokenization\n\nPre-tokenizers split text into pieces before the algorithm runs. Each\npiece carries byte offsets into the original text. The algorithm then\ntokenizes each piece independently.\n\nAvailable pre-tokenizers:\n\n| Pre-tokenizer | Description |\n| --------------------- | --------------------------------------------------------------- |\n| `whitespace ()` | Split on `\\w+\\|[^\\w\\s]+` (word chars grouped, non-word grouped) |\n| `whitespace_split ()` | Split on whitespace (simplest) |\n| `bert ()` | BERT-style: whitespace + punctuation isolation + CJK separation |\n| `byte_level ()` | GPT-2 style byte-level encoding with regex splitting |\n| `punctuation ()` | Separate punctuation from alphanumeric content |\n| `split ~pattern ()` | Split on a literal string pattern |\n| `char_delimiter c` | Split on a single character |\n| `digits ()` | Split on digit boundaries |\n| `metaspace ()` | Replace whitespace with a visible marker (SentencePiece) |\n| `unicode_scripts ()` | Split on Unicode script boundaries |\n| `fixed_length n` | Fixed-size character chunks |\n\nUse `pre_tokenize` to inspect how a pre-tokenizer splits text. It returns\na list of `(piece, (start_offset, end_offset))` pairs:\n\n```ocaml\nopen Brot\n\nlet text = \"Hello, world! How's it going?\"\n\nlet whitespace_pieces =\n Pre_tokenizer.pre_tokenize (Pre_tokenizer.whitespace ()) text\n(* [(\"Hello\", (0,5)); (\",\", (5,6)); (\"world\", (7,12)); (\"!\", (12,13)); ...] *)\n\nlet bert_pieces =\n Pre_tokenizer.pre_tokenize (Pre_tokenizer.bert ()) text\n\nlet punct_pieces =\n Pre_tokenizer.pre_tokenize (Pre_tokenizer.punctuation ()) text\n```\n\nCompose pre-tokenizers with `sequence`. Each pre-tokenizer in the chain\nprocesses the pieces from the previous one:\n\n```ocaml\nopen Brot\n\nlet pre =\n Pre_tokenizer.sequence\n [ Pre_tokenizer.whitespace_split (); Pre_tokenizer.digits () ]\n\nlet pieces = Pre_tokenizer.pre_tokenize pre \"order 42 shipped\"\n(* [(\"order\", _); (\"4\", _); (\"2\", _); (\"shipped\", _)] *)\n```\n\n## Tokenization Algorithms\n\nThe algorithm maps pre-tokenized pieces to token IDs using the vocabulary.\nBrot supports 5 algorithms:\n\n| Algorithm | How it splits | Notable models |\n| --------------- | ------------------------------------------- | ------------------------------ |\n| BPE | Iterative merge of most frequent pairs | GPT-2, GPT-3/4, RoBERTa, LLaMA |\n| WordPiece | Greedy longest-match with `##` prefix | BERT, DistilBERT, Electra |\n| Unigram | Probabilistic segmentation (max likelihood) | T5, ALBERT, mBART, XLNet |\n| Word-level | Whole words, no subword splitting | Simple models, prototyping |\n| Character-level | Each byte is a token | Byte-level fallback |\n\nSee [Choosing an Algorithm](05-algorithms/) for details on each algorithm,\nwhen to use it, and how to configure training.\n\n## Post-processing\n\nPost-processors add special tokens and set type IDs after tokenization.\nThey handle model-specific requirements like `[CLS]`/`[SEP]` for BERT or\n``/` ` for RoBERTa.\n\nAvailable post-processors:\n\n- `bert ~sep ~cls ()` — `[CLS] A [SEP]` or `[CLS] A [SEP] B [SEP]`, type IDs 0/1\n- `roberta ~sep ~cls ()` — ` A ` or ` A B `, all type IDs 0\n- `byte_level ()` — adjust offsets for byte-level encoding\n- `template ~single ()` — custom template with `$A`, `$B`, and literal token placeholders\n- `sequence processors` — chain multiple post-processors\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[UNK]\", 0); (\"[CLS]\", 1); (\"[SEP]\", 2);\n (\"the\", 3); (\"cat\", 4); (\"sat\", 5); (\"how\", 6); (\"are\", 7); (\"you\", 8) ]\n ~specials:(List.map special [ \"[UNK]\"; \"[CLS]\"; \"[SEP]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) ())\n ~decoder:(Decoder.wordpiece ())\n ~unk_token:\"[UNK]\" ()\n\n(* Single sentence: [CLS] the cat sat [SEP] *)\nlet single = encode tokenizer \"the cat sat\"\n\n(* Sentence pair: [CLS] the cat sat [SEP] how are you [SEP] *)\nlet pair = encode tokenizer ~pair:\"how are you\" \"the cat sat\"\n(* type_ids: 0 for first sentence + [CLS]/[SEP], 1 for second + [SEP] *)\nlet type_ids = Encoding.type_ids pair\n```\n\nThe `template` post-processor gives full control over the format. Use `$A`\nand `$B` as sequence placeholders, and literal token names in brackets.\nAppend `:N` to set type IDs:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n word_level\n ~vocab:\n [ (\"[BOS]\", 0); (\"[EOS]\", 1); (\"hello\", 2); (\"world\", 3) ]\n ~specials:(List.map special [ \"[BOS]\"; \"[EOS]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~post:\n (Post_processor.template\n ~single:\"[BOS]:0 $A:0 [EOS]:0\"\n ~pair:\"[BOS]:0 $A:0 [EOS]:0 $B:1 [EOS]:1\"\n ~special_tokens:[ (\"[BOS]\", 0); (\"[EOS]\", 1) ]\n ())\n ~unk_token:\"[UNK]\" ()\n\nlet enc = encode tokenizer \"hello world\"\nlet tokens = Encoding.tokens enc (* [| \"[BOS]\"; \"hello\"; \"world\"; \"[EOS]\" |] *)\nlet type_ids = Encoding.type_ids enc (* [| 0; 0; 0; 0 |] *)\n```\n\n## Decoding\n\nDecoders reverse encoding-specific transformations to produce natural text\nfrom token strings. They operate on token *strings* (looked up from the\nvocabulary), not IDs.\n\nDecoders fall into two categories:\n\n- **Per-token** — transform each token independently: `bpe`, `byte_fallback`, `metaspace`\n- **Collapsing** — process the entire token list as a whole: `byte_level`, `wordpiece`, `replace`, `strip`, `fuse`\n\nThis distinction matters when composing with `sequence`: per-token decoders\npass a list of transformed tokens to the next decoder, while collapsing\ndecoders produce a single result.\n\nAvailable decoders:\n\n| Decoder | Type | Description |\n| ------------------------- | ---------- | ------------------------------------------------ |\n| `bpe ()` | Per-token | Strip end-of-word suffix, insert spaces |\n| `byte_fallback ()` | Per-token | Convert `<0x41>` hex tokens to bytes |\n| `metaspace ()` | Per-token | Convert metaspace markers to spaces |\n| `byte_level ()` | Collapsing | Reverse GPT-2 byte-to-Unicode encoding |\n| `wordpiece ()` | Collapsing | Strip `##` prefix, join subwords |\n| `replace ~pattern ~by ()` | Collapsing | Replace literal pattern in joined text |\n| `strip ()` | Collapsing | Remove leading/trailing characters |\n| `fuse ()` | Collapsing | Concatenate all tokens with no delimiter |\n| `ctc ()` | Per-token | CTC output decoding (deduplication, pad removal) |\n\n```ocaml\nopen Brot\n\n(* WordPiece decoder: strips ## prefix and joins subwords *)\nlet wp = Decoder.wordpiece ()\nlet text = Decoder.decode wp [ \"[CLS]\"; \"play\"; \"##ing\"; \"cat\"; \"##s\"; \"[SEP]\" ]\n(* \"[CLS] playing cats [SEP]\" *)\n\n(* Sequence of decoders *)\nlet seq = Decoder.sequence [ Decoder.fuse (); Decoder.replace ~pattern:\"_\" ~by:\" \" () ]\nlet text2 = Decoder.decode seq [ \"_Hello\"; \"_world\" ]\n(* \" Hello world\" *)\n```\n\nWhen using `Brot.decode`, the tokenizer looks up token strings from the\nvocabulary and then applies the configured decoder automatically.\n\n## Complete Example\n\nHere is a complete BERT-style tokenizer using all 5 pipeline stages:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n (* 1. Normalizer: lowercase and clean text *)\n ~normalizer:(Normalizer.bert ~lowercase:true ())\n (* 2. Pre-tokenizer: BERT-style splitting *)\n ~pre:(Pre_tokenizer.bert ())\n (* 3. Algorithm: WordPiece with ## prefix *)\n ~vocab:\n [ (\"[PAD]\", 0); (\"[UNK]\", 1); (\"[CLS]\", 2); (\"[SEP]\", 3);\n (\"the\", 4); (\"cat\", 5); (\"sat\", 6); (\"on\", 7); (\"mat\", 8);\n (\"play\", 9); (\"##ing\", 10); (\"##ed\", 11); (\"a\", 12) ]\n ~specials:(List.map special [ \"[PAD]\"; \"[UNK]\"; \"[CLS]\"; \"[SEP]\" ])\n ~unk_token:\"[UNK]\" ~pad_token:\"[PAD]\"\n (* 4. Post-processor: add [CLS] and [SEP] *)\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 2) ~sep:(\"[SEP]\", 3) ())\n (* 5. Decoder: strip ## and join *)\n ~decoder:(Decoder.wordpiece ())\n ()\n\n(* \"The Cat\" is normalized to \"the cat\" before tokenization *)\nlet enc = encode tokenizer \"The Cat Played On A Mat\"\nlet tokens = Encoding.tokens enc\n(* [| \"[CLS]\"; \"the\"; \"cat\"; \"play\"; \"##ed\"; \"on\"; \"a\"; \"mat\"; \"[SEP]\" |] *)\n\n(* Decode back, skipping special tokens *)\nlet text = decode tokenizer ~skip_special_tokens:true (Encoding.ids enc)\n(* \"the cat played on a mat\" *)\n```\n"
},
{
"path": "packages/brot/doc/03-pretrained.md",
"content": "# Pretrained Tokenizers\n\nMost users start by loading an existing tokenizer rather than building one\nfrom scratch. Brot reads and writes HuggingFace `tokenizer.json` files and\nseparate vocabulary/merges model files.\n\n## Loading from tokenizer.json\n\nHuggingFace models ship a `tokenizer.json` that contains the algorithm,\nvocabulary, merge rules, and full pipeline configuration. Load it with\n`from_file`:\n\n\n```ocaml\nopen Brot\n\nlet tokenizer = from_file \"path/to/tokenizer.json\" |> Result.get_ok\nlet encoding = encode tokenizer \"Hello world!\"\nlet ids = Encoding.ids encoding\n```\n\n`from_file` returns `(t, string) result`. Handle errors explicitly when\nthe file may be missing or malformed:\n\n\n```ocaml\nlet tokenizer =\n match Brot.from_file \"tokenizer.json\" with\n | Ok t -> t\n | Error msg -> failwith msg\n```\n\n## Loading from Model Files\n\nOlder models ship separate `vocab.json` and `merges.txt` files instead\nof a single `tokenizer.json`. Use `from_model_file`:\n\n\n```ocaml\nopen Brot\n\n(* BPE: provide both vocab and merges *)\nlet tokenizer =\n from_model_file ~vocab:\"vocab.json\" ~merges:\"merges.txt\"\n ~pre:(Pre_tokenizer.byte_level ~add_prefix_space:false ())\n ~decoder:(Decoder.byte_level ())\n ()\n\n(* WordPiece: vocab only, no merges *)\nlet tokenizer =\n from_model_file ~vocab:\"vocab.txt\"\n ~pre:(Pre_tokenizer.bert ())\n ~decoder:(Decoder.wordpiece ())\n ()\n```\n\nWhen `merges` is provided, a BPE tokenizer is created. Without it,\nWordPiece is used. The pipeline stages (normalizer, pre-tokenizer,\npost-processor, decoder) must be configured explicitly since model files\ndo not include them.\n\n## Building Known Pipelines\n\nWhen you need full control over the pipeline or want to understand what\neach stage does, build the tokenizer from scratch with an inline\nvocabulary. The following examples show the standard configurations for\nwell-known models.\n\n### BERT (uncased)\n\nBERT uses WordPiece with `##` continuation prefix, BERT normalization\n(lowercase, clean text, CJK padding), BERT pre-tokenization (whitespace +\npunctuation), and `[CLS]`/`[SEP]` post-processing:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[PAD]\", 0); (\"[UNK]\", 1); (\"[CLS]\", 2); (\"[SEP]\", 3);\n (\"the\", 4); (\"cat\", 5); (\"sat\", 6); (\"on\", 7); (\"mat\", 8);\n (\"play\", 9); (\"##ing\", 10); (\"##ed\", 11); (\"a\", 12);\n (\"is\", 13); (\"good\", 14) ]\n ~normalizer:(Normalizer.bert ~lowercase:true ())\n ~pre:(Pre_tokenizer.bert ())\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 2) ~sep:(\"[SEP]\", 3) ())\n ~decoder:(Decoder.wordpiece ())\n ~specials:(List.map special [ \"[PAD]\"; \"[UNK]\"; \"[CLS]\"; \"[SEP]\" ])\n ~unk_token:\"[UNK]\" ~pad_token:\"[PAD]\" ()\n\nlet enc = encode tokenizer \"The Cat Is Playing\"\nlet tokens = Encoding.tokens enc\n(* [| \"[CLS]\"; \"the\"; \"cat\"; \"is\"; \"play\"; \"##ing\"; \"[SEP]\" |] *)\nlet decoded = decode tokenizer ~skip_special_tokens:true (Encoding.ids enc)\n(* \"the cat is playing\" *)\n```\n\n### GPT-2\n\nGPT-2 uses BPE with byte-level pre-tokenization (no information loss,\nhandles any Unicode input) and byte-level decoding:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n bpe\n ~vocab:\n [ (\"H\", 0); (\"e\", 1); (\"l\", 2); (\"o\", 3); (\"Ġ\", 4); (\"w\", 5);\n (\"r\", 6); (\"d\", 7); (\"He\", 8); (\"ll\", 9); (\"llo\", 10);\n (\"Hello\", 11); (\"Ġw\", 12); (\"or\", 13); (\"ld\", 14);\n (\"orld\", 15); (\"Ġworld\", 16) ]\n ~merges:\n [ (\"H\", \"e\"); (\"l\", \"l\"); (\"ll\", \"o\"); (\"He\", \"llo\");\n (\"Ġ\", \"w\"); (\"o\", \"r\"); (\"l\", \"d\"); (\"or\", \"ld\");\n (\"Ġw\", \"orld\") ]\n ~pre:(Pre_tokenizer.byte_level ~add_prefix_space:false ())\n ~decoder:(Decoder.byte_level ())\n ()\n\nlet enc = encode tokenizer \"Hello world\"\nlet tokens = Encoding.tokens enc (* [| \"Hello\"; \"Ġworld\" |] *)\nlet decoded = decode tokenizer (Encoding.ids enc) (* \"Hello world\" *)\n```\n\n### SentencePiece-style (T5, ALBERT)\n\nSentencePiece models use Unigram with metaspace pre-tokenization (spaces\nreplaced by a visible marker) and metaspace decoding:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n unigram\n ~vocab:\n [ (\"\", -1.0); (\"\\xe2\\x96\\x81\", -2.0);\n (\"\\xe2\\x96\\x81the\", -1.5); (\"\\xe2\\x96\\x81cat\", -1.8);\n (\"\\xe2\\x96\\x81is\", -1.6); (\"\\xe2\\x96\\x81play\", -2.0);\n (\"ing\", -2.5); (\"\\xe2\\x96\\x81a\", -1.4); (\"\\xe2\\x96\\x81good\", -2.1) ]\n ~pre:(Pre_tokenizer.metaspace ~replacement:'\\xe2' ())\n ~decoder:(Decoder.metaspace ~replacement:'\\xe2' ())\n ~unk_token:\"\" ()\n\nlet enc = encode tokenizer \"the cat is playing\"\n```\n\n## Saving Tokenizers\n\nSave a tokenizer in HuggingFace format for later use or sharing:\n\n\n```ocaml\n(* Save as tokenizer.json (full pipeline) *)\nBrot.save_pretrained tokenizer ~path:\"./my_tokenizer\"\n\n(* Save just the vocabulary and merges files *)\nlet files = Brot.save_model_files tokenizer ~folder:\"./model\" ()\n\n(* Export BPE merges in tiktoken format *)\nBrot.export_tiktoken tokenizer\n ~merges_path:\"./tiktoken_merges.txt\"\n ~vocab_path:\"./tiktoken_vocab.txt\"\n```\n\n## Training from Scratch\n\nTrain a tokenizer from a text corpus. Configure the full pipeline\nalongside the training parameters:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n train_bpe\n ~vocab_size:120\n ~min_frequency:1\n ~show_progress:false\n ~pre:(Pre_tokenizer.whitespace ())\n ~specials:(List.map special [ \"[PAD]\"; \"[UNK]\" ])\n ~unk_token:\"[UNK]\" ~pad_token:\"[PAD]\"\n (`Seq (List.to_seq\n [ \"The quick brown fox jumps over the lazy dog.\";\n \"Machine learning models need good tokenizers.\";\n \"Subword tokenization handles unknown words gracefully.\";\n \"The fox jumped over the lazy dog again.\";\n \"Tokenizers convert text to numerical representations.\" ]))\n\nlet size = vocab_size tokenizer\nlet enc = encode tokenizer \"The quick fox\"\n```\n\nSee [Choosing an Algorithm](05-algorithms/) for guidance on which algorithm\nto train and how to tune parameters like `vocab_size`, `min_frequency`,\nand algorithm-specific options.\n"
},
{
"path": "packages/brot/doc/04-batch-processing.md",
"content": "# Batch Processing\n\nReal-world usage requires encoding multiple texts into uniform-length\nsequences for model input. This guide covers encoding metadata, sentence\npairs, batch encoding, padding, truncation, and offset alignment.\n\n## Encoding Metadata\n\n`Encoding.t` carries parallel arrays that all share the same length. Each\nfield serves a specific purpose in model input preparation:\n\n| Field | Type | Description |\n| --------------------- | ------------------- | ----------------------------------------------- |\n| `ids` | `int array` | Token IDs for model input |\n| `tokens` | `string array` | String representation of each token |\n| `offsets` | `(int * int) array` | `(start, end)` byte positions in source text |\n| `type_ids` | `int array` | Segment IDs: 0 for sentence A, 1 for sentence B |\n| `attention_mask` | `int array` | 1 for real tokens, 0 for padding |\n| `special_tokens_mask` | `int array` | 1 for special tokens, 0 for content |\n| `word_ids` | `int option array` | Source word index, or `None` for special tokens |\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[UNK]\", 0); (\"[CLS]\", 1); (\"[SEP]\", 2);\n (\"the\", 3); (\"cat\", 4); (\"play\", 5); (\"##ing\", 6) ]\n ~specials:(List.map special [ \"[UNK]\"; \"[CLS]\"; \"[SEP]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) ())\n ~decoder:(Decoder.wordpiece ())\n ~unk_token:\"[UNK]\" ()\n\nlet enc = encode tokenizer \"the cat playing\"\n(* tokens: [| \"[CLS]\"; \"the\"; \"cat\"; \"play\"; \"##ing\"; \"[SEP]\" |] *)\nlet ids = Encoding.ids enc\nlet type_ids = Encoding.type_ids enc\nlet attention_mask = Encoding.attention_mask enc\nlet special_tokens_mask = Encoding.special_tokens_mask enc\nlet offsets = Encoding.offsets enc\nlet word_ids = Encoding.word_ids enc\n(* word_ids: [| None; Some 0; Some 1; Some 2; Some 2; None |]\n \"play\" and \"##ing\" share word index 2 *)\n```\n\n## Sentence Pairs\n\nMany NLP tasks (question answering, natural language inference, sentence\nsimilarity) operate on pairs of sentences. Use `encode ~pair` to encode\nboth sequences together:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[UNK]\", 0); (\"[CLS]\", 1); (\"[SEP]\", 2);\n (\"the\", 3); (\"cat\", 4); (\"sat\", 5); (\"how\", 6);\n (\"are\", 7); (\"you\", 8) ]\n ~specials:(List.map special [ \"[UNK]\"; \"[CLS]\"; \"[SEP]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) ())\n ~decoder:(Decoder.wordpiece ())\n ~unk_token:\"[UNK]\" ()\n\nlet enc = encode tokenizer ~pair:\"how are you\" \"the cat sat\"\n(* tokens: [| \"[CLS]\"; \"the\"; \"cat\"; \"sat\"; \"[SEP]\"; \"how\"; \"are\"; \"you\"; \"[SEP]\" |] *)\nlet type_ids = Encoding.type_ids enc\n(* [| 0; 0; 0; 0; 0; 1; 1; 1; 1 |] *)\n```\n\nType IDs distinguish the two sentences: 0 for the first sequence\n(including `[CLS]` and first `[SEP]`), 1 for the second (including\nfinal `[SEP]`).\n\n## Batch Encoding\n\nEncode multiple texts at once with `encode_batch`, or multiple sentence\npairs with `encode_pairs_batch`:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[UNK]\", 0); (\"[CLS]\", 1); (\"[SEP]\", 2);\n (\"the\", 3); (\"cat\", 4); (\"sat\", 5);\n (\"how\", 6); (\"are\", 7); (\"you\", 8); (\"good\", 9) ]\n ~specials:(List.map special [ \"[UNK]\"; \"[CLS]\"; \"[SEP]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) ())\n ~decoder:(Decoder.wordpiece ())\n ~unk_token:\"[UNK]\" ()\n\n(* Batch of single sentences *)\nlet encodings =\n encode_batch tokenizer [ \"the cat\"; \"the cat sat\"; \"good\" ]\nlet lengths = List.map Encoding.length encodings\n(* [4; 5; 3] — each includes [CLS] and [SEP] *)\n\n(* Batch of sentence pairs *)\nlet pairs =\n encode_pairs_batch tokenizer\n [ (\"the cat sat\", \"how are you\"); (\"good\", \"the cat\") ]\n```\n\n## Padding\n\nModels require uniform sequence lengths within a batch. Padding extends\nshorter sequences with padding tokens. Three strategies are available:\n\n- **`Batch_longest`** — pad to the longest sequence in the batch\n- **`Fixed n`** — pad every sequence to exactly `n` tokens\n- **`To_multiple n`** — pad to the smallest multiple of `n` that fits\n\nPadding tokens have `attention_mask = 0` and `special_tokens_mask = 1`.\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n word_level\n ~vocab:\n [ (\"[PAD]\", 0); (\"[UNK]\", 1); (\"the\", 2); (\"cat\", 3);\n (\"sat\", 4); (\"on\", 5); (\"a\", 6); (\"mat\", 7) ]\n ~specials:(List.map special [ \"[PAD]\"; \"[UNK]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~unk_token:\"[UNK]\" ~pad_token:\"[PAD]\" ()\n\nlet texts = [ \"the cat\"; \"the cat sat on a mat\"; \"cat\" ]\n\n(* Pad to longest in batch — all encodings have length 6 *)\nlet batch1 =\n encode_batch tokenizer ~padding:(padding `Batch_longest) texts\n\n(* Pad to fixed length — all encodings have length 8 *)\nlet batch2 =\n encode_batch tokenizer ~padding:(padding (`Fixed 8)) texts\n\n(* Pad to multiple of 4 — lengths rounded up to nearest multiple *)\nlet batch3 =\n encode_batch tokenizer ~padding:(padding (`To_multiple 4)) texts\n```\n\nBy default, padding is applied to the right. Use `` ~direction:`Left ``\nfor left-padding, which is common for autoregressive generation:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n word_level\n ~vocab:\n [ (\"[PAD]\", 0); (\"[UNK]\", 1); (\"the\", 2); (\"cat\", 3); (\"sat\", 4) ]\n ~specials:(List.map special [ \"[PAD]\"; \"[UNK]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~unk_token:\"[UNK]\" ~pad_token:\"[PAD]\" ()\n\nlet encodings =\n encode_batch tokenizer\n ~padding:(padding ~direction:`Left (`Fixed 5))\n [ \"the cat\"; \"the cat sat\" ]\n(* tokens: [| \"[PAD]\"; \"[PAD]\"; \"[PAD]\"; \"the\"; \"cat\" |]\n [| \"[PAD]\"; \"[PAD]\"; \"the\"; \"cat\"; \"sat\" |] *)\n```\n\n## Truncation\n\nTruncation limits sequences to a maximum length. Excess tokens are\ntrimmed from the specified direction:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n word_level\n ~vocab:\n [ (\"[UNK]\", 0); (\"the\", 1); (\"quick\", 2); (\"brown\", 3);\n (\"fox\", 4); (\"jumps\", 5); (\"over\", 6) ]\n ~specials:(List.map special [ \"[UNK]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~unk_token:\"[UNK]\" ()\n\nlet text = \"the quick brown fox jumps over\"\n\n(* Truncate from the right (default) *)\nlet enc_right = encode tokenizer ~truncation:(truncation 4) text\nlet tokens_right = Encoding.tokens enc_right\n(* [| \"the\"; \"quick\"; \"brown\"; \"fox\" |] *)\n\n(* Truncate from the left *)\nlet enc_left =\n encode tokenizer ~truncation:(truncation ~direction:`Left 4) text\nlet tokens_left = Encoding.tokens enc_left\n(* [| \"brown\"; \"fox\"; \"jumps\"; \"over\" |] *)\n```\n\nWhen using a post-processor that adds special tokens, account for the\ntokens it adds. Use `Post_processor.added_tokens` to calculate the\nbudget:\n\n```ocaml\nopen Brot\n\nlet post = Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) ()\nlet added_single = Post_processor.added_tokens post ~is_pair:false (* 2 *)\nlet added_pair = Post_processor.added_tokens post ~is_pair:true (* 3 *)\n```\n\n## Padding and Truncation Together\n\nThe common pattern for model input: truncate long sequences and pad short\nones to a uniform length:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n word_level\n ~vocab:\n [ (\"[PAD]\", 0); (\"[UNK]\", 1); (\"the\", 2); (\"cat\", 3);\n (\"sat\", 4); (\"on\", 5); (\"a\", 6); (\"mat\", 7);\n (\"dog\", 8); (\"ran\", 9); (\"fast\", 10) ]\n ~specials:(List.map special [ \"[PAD]\"; \"[UNK]\" ])\n ~pre:(Pre_tokenizer.whitespace ())\n ~unk_token:\"[UNK]\" ~pad_token:\"[PAD]\" ()\n\nlet encodings =\n encode_batch tokenizer\n ~truncation:(truncation 4)\n ~padding:(padding (`Fixed 4))\n [ \"the cat sat on a mat\"; \"the dog ran\"; \"cat\" ]\n(* All encodings have exactly 4 tokens.\n Long sequences are truncated, short ones are padded.\n attention_mask distinguishes real tokens (1) from padding (0). *)\nlet masks = List.map Encoding.attention_mask encodings\n```\n\n## Offsets and Alignment\n\n`Encoding.offsets` maps each token back to its `(start, end)` byte span\nin the original text. This is useful for tasks like named entity\nrecognition where you need to extract the source text for each token:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[UNK]\", 0); (\"hello\", 1); (\"world\", 2);\n (\"play\", 3); (\"##ing\", 4) ]\n ~pre:(Pre_tokenizer.whitespace ())\n ~decoder:(Decoder.wordpiece ())\n ~unk_token:\"[UNK]\" ()\n\nlet text = \"hello playing world\"\nlet enc = encode tokenizer text\nlet offsets = Encoding.offsets enc\n(* offsets.(0) = (0, 5) -> \"hello\"\n offsets.(1) = (6, 13) -> \"playing\" (start of \"play\")\n offsets.(2) = (6, 13) -> \"playing\" (extent of \"##ing\")\n offsets.(3) = (14, 19) -> \"world\" *)\n\n(* Extract source span for a token *)\nlet start, end_ = offsets.(0)\nlet source = String.sub text start (end_ - start) (* \"hello\" *)\n```\n\n`Encoding.word_ids` groups subword tokens back to their source word.\nTokens that belong to the same word share the same word index:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[UNK]\", 0); (\"the\", 1); (\"cat\", 2);\n (\"play\", 3); (\"##ing\", 4); (\"##s\", 5) ]\n ~pre:(Pre_tokenizer.whitespace ())\n ~decoder:(Decoder.wordpiece ())\n ~unk_token:\"[UNK]\" ()\n\nlet enc = encode tokenizer \"the cat playing\"\nlet word_ids = Encoding.word_ids enc\n(* [| Some 0; Some 1; Some 2; Some 2 |]\n \"play\" and \"##ing\" share word index 2,\n indicating they come from the same source word *)\n```\n"
},
{
"path": "packages/brot/doc/05-algorithms.md",
"content": "# Choosing a Tokenization Algorithm\n\nBrot supports 5 tokenization algorithms. The three subword algorithms\n(BPE, WordPiece, Unigram) handle open vocabulary by splitting rare words\ninto smaller pieces. Word-level and character-level are simpler\nalternatives.\n\n## BPE (Byte Pair Encoding)\n\nBPE starts with individual characters and iteratively merges the most\nfrequent adjacent pairs. The merge rules, learned during training, define\nhow text is split. Used by GPT-2, GPT-3/4, RoBERTa, and LLaMA.\n\nConstructor: `Brot.bpe`. Trainer: `Brot.train_bpe`.\n\nKey parameters:\n- `vocab_size` — target vocabulary size (default: 30000)\n- `min_frequency` — minimum pair frequency for merging (default: 0)\n- `dropout` — probability of skipping merges for data augmentation\n- `byte_fallback` — use `<0x00>` byte tokens instead of unknown token\n- `continuing_subword_prefix` — prefix for non-initial subwords\n- `end_of_word_suffix` — suffix marking word boundaries (e.g., ``)\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n bpe\n ~vocab:\n [ (\"h\", 0); (\"e\", 1); (\"l\", 2); (\"o\", 3); (\" \", 4); (\"w\", 5);\n (\"r\", 6); (\"d\", 7); (\"he\", 8); (\"ll\", 9); (\"llo\", 10);\n (\"hello\", 11); (\"wo\", 12); (\"rl\", 13); (\"rld\", 14); (\"world\", 15) ]\n ~merges:\n [ (\"h\", \"e\"); (\"l\", \"l\"); (\"ll\", \"o\"); (\"he\", \"llo\");\n (\"w\", \"o\"); (\"r\", \"l\"); (\"rl\", \"d\"); (\"wo\", \"rld\") ]\n ()\n\nlet enc = encode tokenizer \"hello world\"\nlet tokens = Encoding.tokens enc (* [| \"hello\"; \" \"; \"world\" |] *)\n```\n\nTraining BPE:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n train_bpe ~vocab_size:80 ~min_frequency:1 ~show_progress:false\n (`Seq (List.to_seq\n [ \"The quick brown fox jumps over the lazy dog\";\n \"The dog barked at the brown fox\";\n \"Quick brown foxes are rare and beautiful\" ]))\n\nlet size = vocab_size tokenizer\nlet enc = encode tokenizer \"The brown fox\"\n```\n\n## WordPiece\n\nWordPiece uses a greedy longest-match-first algorithm. For each word, it\nfinds the longest prefix in the vocabulary, then continues with the\nremainder prefixed by a continuation marker (default: `##`). Used by BERT,\nDistilBERT, and Electra.\n\nConstructor: `Brot.wordpiece`. Trainer: `Brot.train_wordpiece`.\n\nKey parameters:\n- `vocab_size` — target vocabulary size (default: 30000)\n- `continuing_subword_prefix` — prefix for non-initial subwords (default: `##`)\n- `max_input_chars_per_word` — words longer than this become unknown (default: 100)\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n wordpiece\n ~vocab:\n [ (\"[UNK]\", 0); (\"the\", 1); (\"cat\", 2); (\"play\", 3);\n (\"##ing\", 4); (\"##ed\", 5); (\"##s\", 6); (\"un\", 7);\n (\"##happy\", 8); (\"##ly\", 9) ]\n ~pre:(Pre_tokenizer.whitespace ())\n ~decoder:(Decoder.wordpiece ())\n ~unk_token:\"[UNK]\" ()\n\nlet enc = encode tokenizer \"the cat playing unhappily\"\nlet tokens = Encoding.tokens enc\n(* [| \"the\"; \"cat\"; \"play\"; \"##ing\"; \"un\"; \"##happy\"; \"##ly\" |] *)\nlet decoded = decode tokenizer (Encoding.ids enc)\n(* \"the cat playing unhappily\" *)\n```\n\nTraining WordPiece:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n train_wordpiece ~vocab_size:80 ~show_progress:false\n (`Seq (List.to_seq\n [ \"The quick brown fox jumps over the lazy dog\";\n \"The dog barked at the brown fox\";\n \"Quick brown foxes are rare and beautiful\" ]))\n\nlet size = vocab_size tokenizer\nlet enc = encode tokenizer \"The brown fox\"\n```\n\n## Unigram\n\nUnigram uses probabilistic segmentation: given a vocabulary of subwords\nwith log-probabilities, it finds the segmentation that maximizes the\ntotal likelihood. Training uses the EM algorithm to iteratively prune the\nvocabulary. Used by T5, ALBERT, mBART, and XLNet.\n\nConstructor: `Brot.unigram`. Trainer: `Brot.train_unigram`.\n\nKey parameters:\n- `vocab_size` — target vocabulary size (default: 8000)\n- `shrinking_factor` — fraction of vocabulary to retain per pruning round (default: 0.75)\n- `max_piece_length` — maximum subword length (default: 16)\n- `n_sub_iterations` — EM sub-iterations per pruning round (default: 2)\n\nVocabulary entries are `(token, score)` pairs where scores are negative\nlog probabilities:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n unigram\n ~vocab:\n [ (\"\", 0.0); (\"the\", -1.0); (\"cat\", -1.5);\n (\"th\", -2.0); (\"e\", -2.5); (\"c\", -3.0); (\"a\", -3.0);\n (\"t\", -3.0); (\"at\", -2.0); (\"he\", -2.0);\n (\"sat\", -1.8); (\"on\", -1.5) ]\n ~unk_token:\"\" ()\n\nlet enc = encode tokenizer \"the cat sat on\"\n```\n\nTraining Unigram:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n train_unigram ~vocab_size:60 ~show_progress:false\n (`Seq (List.to_seq\n [ \"The quick brown fox jumps over the lazy dog\";\n \"The dog barked at the brown fox\";\n \"Quick brown foxes are rare and beautiful\" ]))\n\nlet size = vocab_size tokenizer\nlet enc = encode tokenizer \"The brown fox\"\n```\n\n## Word-level\n\nWord-level tokenization maps each word directly to a token ID. No\nsubword splitting is performed — words not in the vocabulary are replaced\nby the unknown token.\n\nConstructor: `Brot.word_level`. Trainer: `Brot.train_wordlevel`.\n\nBest suited for small controlled vocabularies and prototyping. For\nproduction use with open vocabulary, prefer a subword algorithm.\n\nWhen no pre-tokenizer is specified, `word_level` defaults to\n`Pre_tokenizer.whitespace`.\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n word_level\n ~vocab:\n [ (\"[UNK]\", 0); (\"the\", 1); (\"cat\", 2); (\"sat\", 3);\n (\"on\", 4); (\"a\", 5); (\"mat\", 6) ]\n ~unk_token:\"[UNK]\" ()\n\n(* Known words get their IDs, unknown words become [UNK] *)\nlet enc = encode tokenizer \"the cat sat on a rug\"\nlet tokens = Encoding.tokens enc\n(* [| \"the\"; \"cat\"; \"sat\"; \"on\"; \"a\"; \"[UNK]\" |] *)\nlet ids = Encoding.ids enc\n(* [| 1; 2; 3; 4; 5; 0 |] *)\n```\n\n## Character-level\n\nCharacter-level tokenization maps each byte to a token with ID equal to\nits ordinal value. No vocabulary or training is needed.\n\nConstructor: `Brot.chars`.\n\nUseful as a byte-level fallback or for models that operate directly on\ncharacters:\n\n```ocaml\nopen Brot\n\nlet tokenizer = chars ()\n\nlet enc = encode tokenizer \"Hi!\"\nlet tokens = Encoding.tokens enc (* [| \"H\"; \"i\"; \"!\" |] *)\nlet ids = Encoding.ids enc (* [| 72; 105; 33 |] *)\n```\n\n## Quick Reference\n\n| Algorithm | Splitting strategy | Typical vocab | Notable models | Constructor | Trainer |\n| --------------- | ----------------------------------------- | ------------- | ------------------------- | ------------ | ----------------- |\n| BPE | Iterative merge of frequent pairs | 30K-50K | GPT-2, RoBERTa, LLaMA | `bpe` | `train_bpe` |\n| WordPiece | Greedy longest-match with `##` prefix | 30K | BERT, DistilBERT, Electra | `wordpiece` | `train_wordpiece` |\n| Unigram | Probabilistic max-likelihood segmentation | 8K-32K | T5, ALBERT, mBART, XLNet | `unigram` | `train_unigram` |\n| Word-level | Whole words, no splitting | Varies | Simple models | `word_level` | `train_wordlevel` |\n| Character-level | Each byte is a token | 256 | Byte-level models | `chars` | — |\n"
},
{
"path": "packages/brot/doc/06-hf-tokenizers-comparison.md",
"content": "# Brot vs. HuggingFace Tokenizers -- A Practical Comparison\n\nThis guide explains how Brot relates to Python's [HuggingFace Tokenizers](https://github.com/huggingface/tokenizers), focusing on:\n\n* How core concepts map (tokenizer types, pipeline stages, encoding results)\n* Where the APIs feel similar vs. deliberately different\n* How to translate common HuggingFace patterns into Brot\n\nIf you already use HuggingFace Tokenizers, this should be enough to become productive in Brot quickly.\n\n---\n\n## 1. Big-Picture Differences\n\n| Aspect | HuggingFace Tokenizers (Python) | Brot (OCaml) |\n| ------------------ | ---------------------------------------------------- | ----------------------------------------------------------------------------- |\n| Language | Python bindings over Rust | Native OCaml |\n| Core type | `tokenizers.Tokenizer` | `Brot.t` |\n| Encoding result | `tokenizers.Encoding` | `Encoding.t` |\n| Algorithms | `BPE`, `WordPiece`, `Unigram`, `WordLevel` | `Brot.bpe`, `Brot.wordpiece`, `Brot.unigram`, `Brot.word_level`, `Brot.chars` |\n| Pipeline stages | Mutable properties on `Tokenizer` object | Immutable `~normalizer`, `~pre`, `~post`, `~decoder` args |\n| Mutability | Tokenizer is mutable (set properties after creation) | Tokenizer is immutable after creation |\n| HuggingFace compat | Native format | Full `tokenizer.json` read/write via `from_file`/`save_pretrained` |\n| Training | `Trainer` objects passed to `tokenizer.train()` | `Brot.train_bpe`, `Brot.train_wordpiece`, etc. |\n| Padding config | `tokenizer.enable_padding()` | `~padding` arg on `encode`/`encode_batch` |\n| Truncation config | `tokenizer.enable_truncation()` | `~truncation` arg on `encode`/`encode_batch` |\n\n**Brot semantics to know (read once):**\n- Tokenizers are immutable. Pipeline components are set at construction time, not mutated after.\n- `from_file` returns `(t, string) result`. Handle errors explicitly.\n- Padding and truncation are per-call parameters, not global tokenizer state.\n- Special tokens use a record type (`Brot.special`) with explicit control over stripping and normalization.\n- `encode` returns `Encoding.t`; use `encode_ids` when you only need the ID array.\n\n---\n\n## 2. Loading Pretrained Tokenizers\n\n### 2.1 From a tokenizer.json file\n\n**HuggingFace**\n\n```python\nfrom tokenizers import Tokenizer\n\ntokenizer = Tokenizer.from_file(\"tokenizer.json\")\n```\n\n**Brot**\n\n\n```ocaml\nlet tokenizer = Brot.from_file \"tokenizer.json\" |> Result.get_ok\n```\n\nBoth read the same `tokenizer.json` format. Brot's `from_file` returns a `result` instead of raising an exception.\n\n### 2.2 From vocabulary and merges files\n\n**HuggingFace**\n\n```python\nfrom tokenizers import Tokenizer\nfrom tokenizers.models import BPE\n\ntokenizer = Tokenizer(BPE.from_file(\"vocab.json\", \"merges.txt\"))\n```\n\n**Brot**\n\n\n```ocaml\nlet tokenizer =\n Brot.from_model_file\n ~vocab:\"vocab.json\"\n ~merges:\"merges.txt\"\n ()\n```\n\nWhen `~merges` is omitted, Brot infers WordPiece instead of BPE.\n\n### 2.3 Saving\n\n**HuggingFace**\n\n```python\ntokenizer.save(\"tokenizer.json\")\n```\n\n**Brot**\n\n\n```ocaml\nBrot.save_pretrained tokenizer ~path:\"./my_tokenizer\"\n```\n\n`save_pretrained` creates `path/tokenizer.json` in HuggingFace format. Use `to_json` when you need the JSON value directly.\n\n---\n\n## 3. Encoding Text\n\n### 3.1 Basic encoding\n\n**HuggingFace**\n\n```python\noutput = tokenizer.encode(\"Hello world!\")\noutput.ids # [101, 7592, 2088, 999, 102]\noutput.tokens # ['[CLS]', 'hello', 'world', '!', '[SEP]']\noutput.offsets # [(0, 0), (0, 5), (6, 11), (11, 12), (0, 0)]\noutput.type_ids # [0, 0, 0, 0, 0]\noutput.attention_mask # [1, 1, 1, 1, 1]\n```\n\n**Brot**\n\n\n```ocaml\nlet enc = Brot.encode tokenizer \"Hello world!\"\nlet ids = Encoding.ids enc (* int array *)\nlet toks = Encoding.tokens enc (* string array *)\nlet offs = Encoding.offsets enc (* (int * int) array *)\nlet types = Encoding.type_ids enc (* int array *)\nlet mask = Encoding.attention_mask enc (* int array *)\n```\n\n### 3.2 IDs only\n\n**HuggingFace**\n\n```python\nids = tokenizer.encode(\"Hello world!\").ids\n```\n\n**Brot**\n\n\n```ocaml\nlet ids = Brot.encode_ids tokenizer \"Hello world!\"\n```\n\n`encode_ids` is a shortcut that avoids constructing the full `Encoding.t` when you only need token IDs.\n\n### 3.3 Without special tokens\n\n**HuggingFace**\n\n```python\noutput = tokenizer.encode(\"Hello world!\", add_special_tokens=False)\n```\n\n**Brot**\n\n\n```ocaml\nlet enc = Brot.encode tokenizer ~add_special_tokens:false \"Hello world!\"\n```\n\n---\n\n## 4. Decoding\n\n### 4.1 Basic decoding\n\n**HuggingFace**\n\n```python\ntext = tokenizer.decode([101, 7592, 2088, 999, 102])\ntext_clean = tokenizer.decode([101, 7592, 2088, 999, 102], skip_special_tokens=True)\n```\n\n**Brot**\n\n\n```ocaml\nlet text = Brot.decode tokenizer [| 101; 7592; 2088; 999; 102 |]\nlet text_clean =\n Brot.decode tokenizer ~skip_special_tokens:true\n [| 101; 7592; 2088; 999; 102 |]\n```\n\n### 4.2 Batch decoding\n\n**HuggingFace**\n\n```python\ntexts = tokenizer.decode_batch([[101, 7592, 102], [101, 2088, 102]])\n```\n\n**Brot**\n\n\n```ocaml\nlet texts =\n Brot.decode_batch tokenizer\n [ [| 101; 7592; 102 |]; [| 101; 2088; 102 |] ]\n```\n\n---\n\n## 5. Batch Encoding\n\n**HuggingFace**\n\n```python\noutputs = tokenizer.encode_batch([\"Hello world!\", \"How are you?\"])\n# outputs is a list of Encoding objects\nfor enc in outputs:\n print(enc.ids)\n```\n\n**Brot**\n\n\n```ocaml\nlet encodings =\n Brot.encode_batch tokenizer\n [ \"Hello world!\"; \"How are you?\" ]\n\nlet () =\n List.iter\n (fun enc ->\n let ids = Encoding.ids enc in\n Array.iter (Printf.printf \"%d \") ids;\n print_newline ())\n encodings\n```\n\nBoth return a list of encoding objects, one per input.\n\n---\n\n## 6. Padding and Truncation\n\n### 6.1 Padding\n\nIn HuggingFace, padding is global state on the tokenizer. In Brot, it is a per-call parameter.\n\n**HuggingFace**\n\n```python\ntokenizer.enable_padding(\n direction=\"right\",\n pad_id=0,\n pad_token=\"[PAD]\",\n length=128, # fixed length\n)\noutput = tokenizer.encode(\"Hello\")\n# output.attention_mask shows 0s for padding positions\n```\n\n**Brot**\n\n\n```ocaml\nlet pad = Brot.padding ~pad_id:0 ~pad_token:\"[PAD]\" (`Fixed 128)\nlet enc = Brot.encode tokenizer ~padding:pad \"Hello\"\n(* Encoding.attention_mask enc has 0s for padding positions *)\n```\n\nPadding strategies:\n\n| HuggingFace | Brot |\n| --------------------------------------- | ----------------------------- |\n| `length=None` (pad to longest in batch) | `` `Batch_longest `` |\n| `length=128` (fixed) | `` `Fixed 128 `` |\n| `pad_to_multiple_of=8` | `` `To_multiple 8 `` |\n| `direction=\"left\"` | `~direction:`Left` |\n| `direction=\"right\"` (default) | `~direction:`Right` (default) |\n\n### 6.2 Truncation\n\n**HuggingFace**\n\n```python\ntokenizer.enable_truncation(max_length=512, direction=\"right\")\noutput = tokenizer.encode(\"Very long text ...\")\n```\n\n**Brot**\n\n\n```ocaml\nlet trunc = Brot.truncation 512\nlet enc = Brot.encode tokenizer ~truncation:trunc \"Very long text ...\"\n```\n\nTruncation direction defaults to `` `Right `` in both libraries.\n\n### 6.3 Combined padding and truncation\n\n**HuggingFace**\n\n```python\ntokenizer.enable_padding(length=512, pad_token=\"[PAD]\", pad_id=0)\ntokenizer.enable_truncation(max_length=512)\noutputs = tokenizer.encode_batch(texts)\n```\n\n**Brot**\n\n\n```ocaml\nlet pad = Brot.padding ~pad_token:\"[PAD]\" ~pad_id:0 (`Fixed 512)\nlet trunc = Brot.truncation 512\nlet encodings =\n Brot.encode_batch tokenizer ~padding:pad ~truncation:trunc texts\n```\n\nThe key difference: Brot passes these as arguments, so different calls can use different settings without mutating the tokenizer.\n\n---\n\n## 7. Sentence Pairs\n\n**HuggingFace**\n\n```python\n# Single pair\noutput = tokenizer.encode(\"premise\", \"hypothesis\")\noutput.type_ids # [0, 0, 0, 0, 1, 1, 1] (with BERT post-processor)\n\n# Batch of pairs\noutputs = tokenizer.encode_batch([(\"premise1\", \"hyp1\"), (\"premise2\", \"hyp2\")])\n```\n\n**Brot**\n\n\n```ocaml\n(* Single pair *)\nlet enc = Brot.encode tokenizer ~pair:\"hypothesis\" \"premise\"\nlet type_ids = Encoding.type_ids enc (* 0s for first, 1s for second *)\n\n(* Batch of pairs *)\nlet encodings =\n Brot.encode_pairs_batch tokenizer\n [ (\"premise1\", \"hyp1\"); (\"premise2\", \"hyp2\") ]\n```\n\nBrot uses the `~pair` optional argument on `encode` for single pairs and a dedicated `encode_pairs_batch` for batches, instead of overloading the same function with tuples.\n\n---\n\n## 8. Special Tokens\n\n### 8.1 Defining special tokens\n\n**HuggingFace**\n\n```python\nfrom tokenizers import AddedToken\n\ntokenizer.add_special_tokens([\n AddedToken(\"[CLS]\", single_word=False, lstrip=False, rstrip=False),\n AddedToken(\"[SEP]\", single_word=False, lstrip=False, rstrip=False),\n AddedToken(\"[PAD]\", single_word=False, lstrip=False, rstrip=False),\n])\n```\n\n**Brot**\n\n\n```ocaml\nlet tokenizer =\n Brot.bpe\n ~specials:[\n Brot.special \"[CLS]\";\n Brot.special \"[SEP]\";\n Brot.special \"[PAD]\";\n ]\n ~pad_token:\"[PAD]\"\n ~bos_token:\"[CLS]\"\n ~eos_token:\"[SEP]\"\n ()\n```\n\nIn HuggingFace, special tokens are added after construction. In Brot, they are part of construction since tokenizers are immutable. The `special` function accepts optional `~single_word`, `~lstrip`, `~rstrip`, and `~normalized` parameters matching `AddedToken`.\n\n### 8.2 Role tokens\n\n**HuggingFace**\n\n```python\ntokenizer.pad_token # \"[PAD]\"\ntokenizer.cls_token # \"[CLS]\"\ntokenizer.sep_token # \"[SEP]\"\ntokenizer.unk_token # \"[UNK]\"\n```\n\n**Brot**\n\n\n```ocaml\nlet pad = Brot.pad_token tokenizer (* string option *)\nlet bos = Brot.bos_token tokenizer (* string option *)\nlet eos = Brot.eos_token tokenizer (* string option *)\nlet unk = Brot.unk_token tokenizer (* string option *)\n```\n\nBrot uses `bos_token`/`eos_token` instead of `cls_token`/`sep_token` since these are model-agnostic roles. They return `option` instead of raising on missing tokens.\n\n### 8.3 Special tokens mask\n\nBoth libraries provide a mask distinguishing special tokens from content tokens in the encoding:\n\n**HuggingFace**\n\n```python\noutput.special_tokens_mask # [1, 0, 0, 0, 1]\n```\n\n**Brot**\n\n\n```ocaml\nlet mask = Encoding.special_tokens_mask enc (* int array: 1 for special, 0 for content *)\n```\n\n---\n\n## 9. Pipeline Components\n\nBoth libraries use the same four-stage pipeline: normalizer, pre-tokenizer, post-processor, decoder. The difference is how they are configured.\n\n### 9.1 Normalizer\n\n**HuggingFace**\n\n```python\nfrom tokenizers import normalizers\n\ntokenizer.normalizer = normalizers.Sequence([\n normalizers.NFD(),\n normalizers.StripAccents(),\n normalizers.Lowercase(),\n])\n```\n\n**Brot**\n\n\n```ocaml\nlet norm =\n Normalizer.sequence\n [ Normalizer.nfd; Normalizer.strip_accents; Normalizer.lowercase ]\n\nlet tokenizer = Brot.bpe ~normalizer:norm ()\n```\n\nCommon normalizers:\n\n| HuggingFace | Brot |\n| ----------------------------------- | ------------------------------------------ |\n| `normalizers.NFC()` | `Normalizer.nfc` |\n| `normalizers.NFD()` | `Normalizer.nfd` |\n| `normalizers.NFKC()` | `Normalizer.nfkc` |\n| `normalizers.NFKD()` | `Normalizer.nfkd` |\n| `normalizers.Lowercase()` | `Normalizer.lowercase` |\n| `normalizers.StripAccents()` | `Normalizer.strip_accents` |\n| `normalizers.Strip()` | `Normalizer.strip ()` |\n| `normalizers.Replace(pattern, rep)` | `Normalizer.replace ~pattern ~replacement` |\n| `normalizers.Prepend(s)` | `Normalizer.prepend s` |\n| `normalizers.BertNormalizer()` | `Normalizer.bert ()` |\n| `normalizers.ByteLevel()` | `Normalizer.byte_level ()` |\n| `normalizers.Sequence([...])` | `Normalizer.sequence [...]` |\n\n### 9.2 Pre-tokenizer\n\n**HuggingFace**\n\n```python\nfrom tokenizers import pre_tokenizers\n\ntokenizer.pre_tokenizer = pre_tokenizers.Sequence([\n pre_tokenizers.WhitespaceSplit(),\n pre_tokenizers.Punctuation(),\n])\n```\n\n**Brot**\n\n\n```ocaml\nlet pre =\n Pre_tokenizer.sequence\n [ Pre_tokenizer.whitespace_split ();\n Pre_tokenizer.punctuation () ]\n\nlet tokenizer = Brot.bpe ~pre ()\n```\n\nCommon pre-tokenizers:\n\n| HuggingFace | Brot |\n| -------------------------------------- | ----------------------------------- |\n| `pre_tokenizers.Whitespace()` | `Pre_tokenizer.whitespace ()` |\n| `pre_tokenizers.WhitespaceSplit()` | `Pre_tokenizer.whitespace_split ()` |\n| `pre_tokenizers.BertPreTokenizer()` | `Pre_tokenizer.bert ()` |\n| `pre_tokenizers.ByteLevel()` | `Pre_tokenizer.byte_level ()` |\n| `pre_tokenizers.Punctuation()` | `Pre_tokenizer.punctuation ()` |\n| `pre_tokenizers.Digits()` | `Pre_tokenizer.digits ()` |\n| `pre_tokenizers.Metaspace()` | `Pre_tokenizer.metaspace ()` |\n| `pre_tokenizers.UnicodeScripts()` | `Pre_tokenizer.unicode_scripts ()` |\n| `pre_tokenizers.CharDelimiterSplit(c)` | `Pre_tokenizer.char_delimiter c` |\n| `pre_tokenizers.Split(pattern, ...)` | `Pre_tokenizer.split ~pattern ()` |\n| `pre_tokenizers.Sequence([...])` | `Pre_tokenizer.sequence [...]` |\n\n### 9.3 Post-processor\n\n**HuggingFace**\n\n```python\nfrom tokenizers import processors\n\ntokenizer.post_processor = processors.BertProcessing(\n sep=(\"[SEP]\", 102),\n cls=(\"[CLS]\", 101),\n)\n```\n\n**Brot**\n\n\n```ocaml\nlet post =\n Post_processor.bert\n ~sep:(\"[SEP]\", 102)\n ~cls:(\"[CLS]\", 101)\n ()\n\nlet tokenizer = Brot.bpe ~post ()\n```\n\nCommon post-processors:\n\n| HuggingFace | Brot |\n| ------------------------------------------------------------- | ---------------------------------------------------------- |\n| `processors.BertProcessing(sep, cls)` | `Post_processor.bert ~sep ~cls ()` |\n| `processors.RobertaProcessing(sep, cls)` | `Post_processor.roberta ~sep ~cls ()` |\n| `processors.ByteLevel()` | `Post_processor.byte_level ()` |\n| `processors.TemplateProcessing(single, pair, special_tokens)` | `Post_processor.template ~single ?pair ~special_tokens ()` |\n| `processors.Sequence([...])` | `Post_processor.sequence [...]` |\n\n### 9.4 Decoder\n\n**HuggingFace**\n\n```python\nfrom tokenizers import decoders\n\ntokenizer.decoder = decoders.WordPiece(prefix=\"##\")\n```\n\n**Brot**\n\n\n```ocaml\nlet dec = Decoder.wordpiece ~prefix:\"##\" ()\nlet tokenizer = Brot.wordpiece ~decoder:dec ()\n```\n\nCommon decoders:\n\n| HuggingFace | Brot |\n| ------------------------------- | --------------------------------- |\n| `decoders.BPEDecoder(suffix)` | `Decoder.bpe ~suffix ()` |\n| `decoders.ByteLevel()` | `Decoder.byte_level ()` |\n| `decoders.ByteFallback()` | `Decoder.byte_fallback ()` |\n| `decoders.WordPiece(prefix)` | `Decoder.wordpiece ~prefix ()` |\n| `decoders.Metaspace()` | `Decoder.metaspace ()` |\n| `decoders.CTC()` | `Decoder.ctc ()` |\n| `decoders.Replace(pattern, by)` | `Decoder.replace ~pattern ~by ()` |\n| `decoders.Strip()` | `Decoder.strip ()` |\n| `decoders.Fuse()` | `Decoder.fuse ()` |\n| `decoders.Sequence([...])` | `Decoder.sequence [...]` |\n\n### 9.5 Inspecting the pipeline\n\n**HuggingFace**\n\n```python\ntokenizer.normalizer\ntokenizer.pre_tokenizer\ntokenizer.post_processor\ntokenizer.decoder\n```\n\n**Brot**\n\n\n```ocaml\nlet norm = Brot.normalizer tokenizer (* Normalizer.t option *)\nlet pre = Brot.pre_tokenizer tokenizer (* Pre_tokenizer.t option *)\nlet post = Brot.post_processor tokenizer (* Post_processor.t option *)\nlet dec = Brot.decoder tokenizer (* Decoder.t option *)\n```\n\nBrot returns `option` for each stage, since any stage can be absent.\n\n---\n\n## 10. Training Tokenizers\n\n### 10.1 BPE training\n\n**HuggingFace**\n\n```python\nfrom tokenizers import Tokenizer\nfrom tokenizers.models import BPE\nfrom tokenizers.trainers import BpeTrainer\n\ntokenizer = Tokenizer(BPE())\ntrainer = BpeTrainer(\n vocab_size=30000,\n min_frequency=2,\n special_tokens=[\"[UNK]\", \"[CLS]\", \"[SEP]\", \"[PAD]\"],\n)\ntokenizer.train([\"corpus.txt\"], trainer)\n```\n\n**Brot**\n\n\n```ocaml\nlet tokenizer =\n Brot.train_bpe\n (`Files [ \"corpus.txt\" ])\n ~vocab_size:30000\n ~min_frequency:2\n ~specials:[\n Brot.special \"[UNK]\";\n Brot.special \"[CLS]\";\n Brot.special \"[SEP]\";\n Brot.special \"[PAD]\";\n ]\n ~unk_token:\"[UNK]\"\n ~pad_token:\"[PAD]\"\n```\n\nBrot combines the `Tokenizer` + `Trainer` pattern into a single function call. Training data is passed as `` `Files `` (file paths) or `` `Seq `` (string sequence).\n\n### 10.2 WordPiece training\n\n**HuggingFace**\n\n```python\nfrom tokenizers.models import WordPiece\nfrom tokenizers.trainers import WordPieceTrainer\n\ntokenizer = Tokenizer(WordPiece(unk_token=\"[UNK]\"))\ntrainer = WordPieceTrainer(vocab_size=30000, special_tokens=[\"[UNK]\", \"[PAD]\"])\ntokenizer.train([\"corpus.txt\"], trainer)\n```\n\n**Brot**\n\n\n```ocaml\nlet tokenizer =\n Brot.train_wordpiece\n (`Files [ \"corpus.txt\" ])\n ~vocab_size:30000\n ~unk_token:\"[UNK]\"\n ~specials:[ Brot.special \"[UNK]\"; Brot.special \"[PAD]\" ]\n ~pad_token:\"[PAD]\"\n```\n\n### 10.3 Unigram training\n\n**HuggingFace**\n\n```python\nfrom tokenizers.models import Unigram\nfrom tokenizers.trainers import UnigramTrainer\n\ntokenizer = Tokenizer(Unigram())\ntrainer = UnigramTrainer(vocab_size=8000, special_tokens=[\"\", \"\"])\ntokenizer.train([\"corpus.txt\"], trainer)\n```\n\n**Brot**\n\n\n```ocaml\nlet tokenizer =\n Brot.train_unigram\n (`Files [ \"corpus.txt\" ])\n ~vocab_size:8000\n ~unk_token:\"\"\n ~specials:[ Brot.special \"\"; Brot.special \"\" ]\n ~pad_token:\"\"\n```\n\n### 10.4 Training from in-memory data\n\n**HuggingFace**\n\n```python\nfrom tokenizers import Tokenizer\nfrom tokenizers.models import BPE\nfrom tokenizers.trainers import BpeTrainer\n\ntokenizer = Tokenizer(BPE())\ntrainer = BpeTrainer(vocab_size=1000)\ntokenizer.train_from_iterator(\n [\"Hello world\", \"How are you?\", \"Hello again\"],\n trainer,\n)\n```\n\n**Brot**\n\n\n```ocaml\nlet texts = [ \"Hello world\"; \"How are you?\"; \"Hello again\" ]\nlet tokenizer =\n Brot.train_bpe (`Seq (List.to_seq texts)) ~vocab_size:1000\n```\n\n### 10.5 Extending an existing tokenizer\n\n**HuggingFace**\n\n```python\n# Load, then retrain with more data\ntokenizer = Tokenizer.from_file(\"tokenizer.json\")\ntrainer = BpeTrainer(vocab_size=50000)\ntokenizer.train([\"more_data.txt\"], trainer)\n```\n\n**Brot**\n\n\n```ocaml\nlet base = Brot.from_file \"tokenizer.json\" |> Result.get_ok\nlet tokenizer =\n Brot.train_bpe ~init:base (`Files [ \"more_data.txt\" ]) ~vocab_size:50000\n```\n\nThe `~init` parameter on training functions lets you extend an existing tokenizer with additional data.\n\n---\n\n## 11. Vocabulary Inspection\n\n**HuggingFace**\n\n```python\ntokenizer.get_vocab() # dict: token -> id\ntokenizer.get_vocab_size() # int\ntokenizer.token_to_id(\"[CLS]\") # int or None\ntokenizer.id_to_token(101) # str or None\n```\n\n**Brot**\n\n\n```ocaml\nlet v = Brot.vocab tokenizer (* (string * int) list *)\nlet size = Brot.vocab_size tokenizer (* int *)\nlet id = Brot.token_to_id tokenizer \"[CLS]\" (* int option *)\nlet token = Brot.id_to_token tokenizer 101 (* string option *)\n```\n\n`vocab` returns an association list instead of a dictionary. `token_to_id` and `id_to_token` return `option` instead of nullable values.\n\n---\n\n## 12. Quick Cheat Sheet\n\n| Task | HuggingFace Tokenizers | Brot |\n| ------------------- | ----------------------------------------------------------- | ---------------------------------------------------------------- |\n| Load from file | `Tokenizer.from_file(\"tokenizer.json\")` | `Brot.from_file \"tokenizer.json\"` |\n| Save to file | `tokenizer.save(\"tokenizer.json\")` | `Brot.save_pretrained tokenizer ~path:\"./out\"` |\n| Encode text | `tokenizer.encode(\"Hello\")` | `Brot.encode tokenizer \"Hello\"` |\n| Encode IDs only | `tokenizer.encode(\"Hello\").ids` | `Brot.encode_ids tokenizer \"Hello\"` |\n| Encode batch | `tokenizer.encode_batch([\"a\", \"b\"])` | `Brot.encode_batch tokenizer [\"a\"; \"b\"]` |\n| Encode pair | `tokenizer.encode(\"a\", \"b\")` | `Brot.encode tokenizer ~pair:\"b\" \"a\"` |\n| Encode pairs batch | `tokenizer.encode_batch([(\"a\",\"b\"), ...])` | `Brot.encode_pairs_batch tokenizer [(\"a\",\"b\"); ...]` |\n| Decode | `tokenizer.decode(ids)` | `Brot.decode tokenizer ids` |\n| Decode batch | `tokenizer.decode_batch([ids1, ids2])` | `Brot.decode_batch tokenizer [ids1; ids2]` |\n| Get token IDs | `output.ids` | `Encoding.ids enc` |\n| Get tokens | `output.tokens` | `Encoding.tokens enc` |\n| Get attention mask | `output.attention_mask` | `Encoding.attention_mask enc` |\n| Get type IDs | `output.type_ids` | `Encoding.type_ids enc` |\n| Get offsets | `output.offsets` | `Encoding.offsets enc` |\n| Padding | `tokenizer.enable_padding(length=128)` | `Brot.encode tokenizer ~padding:(Brot.padding (`Fixed 128)) ...` |\n| Truncation | `tokenizer.enable_truncation(max_length=512)` | `Brot.encode tokenizer ~truncation:(Brot.truncation 512) ...` |\n| Vocab size | `tokenizer.get_vocab_size()` | `Brot.vocab_size tokenizer` |\n| Token to ID | `tokenizer.token_to_id(\"[CLS]\")` | `Brot.token_to_id tokenizer \"[CLS]\"` |\n| ID to token | `tokenizer.id_to_token(101)` | `Brot.id_to_token tokenizer 101` |\n| Train BPE | `tokenizer.train(files, BpeTrainer(...))` | `Brot.train_bpe (`Files files) ~vocab_size:30000` |\n| Train WordPiece | `tokenizer.train(files, WordPieceTrainer(...))` | `Brot.train_wordpiece (`Files files) ~vocab_size:30000` |\n| Train Unigram | `tokenizer.train(files, UnigramTrainer(...))` | `Brot.train_unigram (`Files files) ~vocab_size:8000` |\n| Train from iterator | `tokenizer.train_from_iterator(iter, trainer)` | `Brot.train_bpe (`Seq seq) ~vocab_size:1000` |\n| Set normalizer | `tokenizer.normalizer = normalizers.Lowercase()` | `Brot.bpe ~normalizer:Normalizer.lowercase ()` |\n| Set pre-tokenizer | `tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()` | `Brot.bpe ~pre:(Pre_tokenizer.byte_level ()) ()` |\n| Set post-processor | `tokenizer.post_processor = processors.BertProcessing(...)` | `Brot.bpe ~post:(Post_processor.bert ~sep ~cls ()) ()` |\n| Set decoder | `tokenizer.decoder = decoders.WordPiece()` | `Brot.bpe ~decoder:(Decoder.wordpiece ()) ()` |\n| Add special tokens | `tokenizer.add_special_tokens([AddedToken(...)])` | Pass `~specials:[Brot.special \"...\"; ...]` at construction |\n"
},
{
"path": "packages/brot/doc/dune",
"content": "(mdx\n (files *.md)\n (package brot)\n (libraries brot))\n"
},
{
"path": "packages/brot/doc/index.md",
"content": "# Brot\n\nBrot tokenizes text into token IDs for language models and reverses the\nprocess. It supports BPE, WordPiece, Unigram, word-level, and\ncharacter-level algorithms, loads and saves HuggingFace `tokenizer.json`\nfiles, and is 1.3-6x faster than HuggingFace tokenizers on most\nbenchmarks.\n\n## Features\n\n- **Tokenization algorithms**: BPE, WordPiece, Unigram, word-level, character-level\n- **HuggingFace compatible**: load and save `tokenizer.json`, load vocab/merges model files\n- **Composable pipeline**: normalizer, pre-tokenizer, post-processor, decoder — each stage independently configurable\n- **Rich encoding output**: token IDs, string tokens, byte offsets, attention masks, type IDs, word IDs, special token masks\n- **Training**: train BPE, WordPiece, Unigram, and word-level tokenizers from scratch\n- **Performance**: 1.3-6x faster than HuggingFace tokenizers (Rust native)\n\n## Quick Start\n\nBuild a BPE tokenizer from a vocabulary and merge rules, encode text,\nand decode it back:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n bpe\n ~vocab:\n [ (\"h\", 0); (\"e\", 1); (\"l\", 2); (\"o\", 3); (\" \", 4); (\"w\", 5);\n (\"r\", 6); (\"d\", 7); (\"he\", 8); (\"ll\", 9); (\"llo\", 10);\n (\"hello\", 11); (\"wo\", 12); (\"rl\", 13); (\"rld\", 14); (\"world\", 15) ]\n ~merges:\n [ (\"h\", \"e\"); (\"l\", \"l\"); (\"ll\", \"o\"); (\"he\", \"llo\");\n (\"w\", \"o\"); (\"r\", \"l\"); (\"rl\", \"d\"); (\"wo\", \"rld\") ]\n ()\n\nlet encoding = encode tokenizer \"hello world\"\nlet ids = Encoding.ids encoding (* [| 11; 4; 15 |] *)\nlet tokens = Encoding.tokens encoding (* [| \"hello\"; \" \"; \"world\" |] *)\nlet decoded = decode tokenizer ids (* \"hello world\" *)\n```\n\nLoad a pretrained tokenizer from a HuggingFace `tokenizer.json` file:\n\n\n```ocaml\nopen Brot\n\nlet tokenizer = from_file \"tokenizer.json\" |> Result.get_ok\nlet encoding = encode tokenizer \"Hello world!\"\nlet ids = Encoding.ids encoding\n```\n\nTrain a tokenizer from a text corpus:\n\n```ocaml\nopen Brot\n\nlet tokenizer =\n train_bpe ~vocab_size:100 ~show_progress:false\n (`Seq (List.to_seq\n [ \"The quick brown fox jumps over the lazy dog\";\n \"The dog barked at the fox\";\n \"Quick brown foxes are rare\" ]))\n\nlet size = vocab_size tokenizer\nlet ids = encode_ids tokenizer \"The quick fox\"\n```\n\n## Next Steps\n\n- [Getting Started](01-getting-started/) — encode, decode, pipeline basics, training\n- [The Tokenization Pipeline](02-pipeline/) — how the 5 pipeline stages work\n- [Pretrained Tokenizers](03-pretrained/) — loading, saving, and building known model pipelines\n- [Batch Processing](04-batch-processing/) — padding, truncation, encoding metadata\n- [Choosing an Algorithm](05-algorithms/) — BPE vs WordPiece vs Unigram and when to use each\n"
},
{
"path": "packages/brot/examples/01-encode-decode/README.md",
"content": "# `01-encode-decode`\n\nYour first tokenizer. This example shows the minimal steps to encode text into\ntoken IDs and decode back.\n\n```bash\ndune exec brot/examples/01-encode-decode/main.exe\n```\n\n## What You'll Learn\n\n- Creating a BPE tokenizer with `Brot.bpe`\n- Encoding text with `Brot.encode`\n- Inspecting token strings and IDs with `Encoding.tokens` and `Encoding.ids`\n- Decoding token IDs back to text with `Brot.decode`\n\n## Key Functions\n\n| Function | Purpose |\n| ----------------- | ------------------------------------------------------ |\n| `bpe` | Create a BPE tokenizer from vocabulary and merge rules |\n| `encode` | Encode text into an `Encoding.t` |\n| `Encoding.ids` | Get the integer token IDs |\n| `Encoding.tokens` | Get the string token representations |\n| `decode` | Convert token IDs back to text |\n\n## How BPE Works\n\nBPE (Byte Pair Encoding) iteratively merges the most frequent character pairs.\nGiven the text `\"hello\"` and merge rules like `(\"h\",\"e\")`, `(\"l\",\"l\")`,\n`(\"he\",\"l\")`, `(\"ll\",\"o\")`, `(\"hel\",\"lo\")`, BPE applies merges in priority\norder until no more merges apply, producing `\"hello\"` as a single token.\n\n## Try It\n\n1. Remove some merge rules and run again to see how the text gets split into\n smaller subword pieces.\n2. Add a new word like `\"held\"` to the vocabulary and encode `\"hello held\"`.\n\n## Next Steps\n\nContinue to [02-encoding-fields](../02-encoding-fields/) to learn about all the\nmetadata in an encoding.\n"
},
{
"path": "packages/brot/examples/01-encode-decode/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/01-encode-decode/main.ml",
"content": "(* Encode and decode.\n\n The simplest possible tokenization: convert text to token IDs and back.\n Demonstrates creating a BPE tokenizer from an inline vocabulary and merge\n rules, encoding text, inspecting tokens and IDs, and decoding. *)\n\nopen Brot\n\nlet () =\n (* Build a small BPE tokenizer. The vocabulary maps token strings to IDs.\n Merge rules define which character pairs to combine, in priority order. *)\n let vocab =\n [\n (\"h\", 0);\n (\"e\", 1);\n (\"l\", 2);\n (\"o\", 3);\n (\" \", 4);\n (\"w\", 5);\n (\"r\", 6);\n (\"d\", 7);\n (\"he\", 8);\n (\"ll\", 9);\n (\"llo\", 10);\n (\"hello\", 11);\n (\"wo\", 12);\n (\"rl\", 13);\n (\"rld\", 14);\n (\"world\", 15);\n ]\n in\n let merges =\n [\n (\"h\", \"e\");\n (\"l\", \"l\");\n (\"ll\", \"o\");\n (\"he\", \"llo\");\n (\"w\", \"o\");\n (\"r\", \"l\");\n (\"rl\", \"d\");\n (\"wo\", \"rld\");\n ]\n in\n let tokenizer = bpe ~vocab ~merges () in\n\n (* Encode text into an Encoding *)\n let text = \"hello world\" in\n let encoding = encode tokenizer text in\n let ids = Encoding.ids encoding in\n let tokens = Encoding.tokens encoding in\n\n Printf.printf \"Text: %S\\n\" text;\n Printf.printf \"Tokens: [%s]\\n\"\n (String.concat \"; \"\n (List.map (fun s -> Printf.sprintf \"%S\" s) (Array.to_list tokens)));\n Printf.printf \"IDs: [%s]\\n\"\n (String.concat \"; \" (Array.to_list (Array.map string_of_int ids)));\n\n (* Decode token IDs back to text *)\n let decoded = decode tokenizer ids in\n Printf.printf \"Decoded: %S\\n\\n\" decoded;\n\n Printf.printf \"Round-trip matches: %b\\n\\n\" (String.equal text decoded);\n\n (* Try another text -- unknown characters become individual tokens *)\n let text2 = \"hello\" in\n let enc2 = encode tokenizer text2 in\n Printf.printf \"Text: %S\\n\" text2;\n Printf.printf \"Tokens: [%s]\\n\"\n (String.concat \"; \"\n (List.map\n (fun s -> Printf.sprintf \"%S\" s)\n (Array.to_list (Encoding.tokens enc2))));\n Printf.printf \"IDs: [%s]\\n\"\n (String.concat \"; \"\n (Array.to_list (Array.map string_of_int (Encoding.ids enc2))))\n"
},
{
"path": "packages/brot/examples/02-encoding-fields/README.md",
"content": "# `02-encoding-fields`\n\nUnderstanding encodings. An `Encoding.t` bundles token IDs with alignment\nmetadata: byte offsets, word indices, type IDs, attention masks, and\nspecial-token flags.\n\n```bash\ndune exec brot/examples/02-encoding-fields/main.exe\n```\n\n## What You'll Learn\n\n- All parallel arrays in an `Encoding.t` and how they align\n- Byte offsets that map each token back to the original text\n- Word indices that group subword tokens by source word\n- Attention mask (1 = real token, 0 = padding)\n- Special tokens mask (1 = special, 0 = content)\n\n## Key Functions\n\n| Function | Purpose |\n| ------------------------------ | ------------------------------------------------- |\n| `Encoding.ids` | Token ID array for model input |\n| `Encoding.tokens` | String representation of each token |\n| `Encoding.offsets` | `(start, end)` byte spans in the original text |\n| `Encoding.word_ids` | Source word index per token (`None` for specials) |\n| `Encoding.type_ids` | Segment IDs (0 or 1 for sentence pairs) |\n| `Encoding.attention_mask` | 1 for real tokens, 0 for padding |\n| `Encoding.special_tokens_mask` | 1 for special tokens, 0 for content |\n| `Encoding.length` | Number of tokens |\n\n## Offsets\n\nOffsets are byte positions `(start, end)` into the original text. You can\nextract the original substring with `String.sub text start (end - start)`.\nThis is essential for highlighting, named entity recognition, and other tasks\nthat need to map tokens back to source text.\n\n## Try It\n\n1. Add more words to the vocabulary and encode a longer sentence.\n2. Encode a text with unknown words and observe the `[UNK]` token.\n\n## Next Steps\n\nContinue to [03-normalizers](../03-normalizers/) to learn how text is cleaned\nbefore tokenization.\n"
},
{
"path": "packages/brot/examples/02-encoding-fields/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/02-encoding-fields/main.ml",
"content": "(* Understanding encodings.\n\n An Encoding bundles token IDs with alignment metadata: byte offsets, word\n indices, segment type IDs, attention masks, and special-token flags. All\n arrays share the same length. *)\n\nopen Brot\n\nlet print_encoding enc =\n let ids = Encoding.ids enc in\n let tokens = Encoding.tokens enc in\n let offsets = Encoding.offsets enc in\n let word_ids = Encoding.word_ids enc in\n let type_ids = Encoding.type_ids enc in\n let attn = Encoding.attention_mask enc in\n let special = Encoding.special_tokens_mask enc in\n\n Printf.printf \"%-6s %-10s %-4s %-12s %-8s %-8s %-6s %-8s\\n\" \"Index\" \"Token\"\n \"ID\" \"Offsets\" \"Word_ID\" \"Type_ID\" \"Attn\" \"Special\";\n Printf.printf \"%s\\n\" (String.make 66 '-');\n\n for i = 0 to Encoding.length enc - 1 do\n let s, e = offsets.(i) in\n let word =\n match word_ids.(i) with Some w -> string_of_int w | None -> \"-\"\n in\n Printf.printf \"%-6d %-10s %-4d (%2d, %2d) %-8s %-8d %-6d %-8d\\n\" i\n tokens.(i) ids.(i) s e word type_ids.(i) attn.(i) special.(i)\n done\n\nlet () =\n (* Word-level tokenizer: each word maps to one token *)\n let vocab =\n [\n (\"[UNK]\", 0);\n (\"hello\", 1);\n (\"world\", 2);\n (\"the\", 3);\n (\"is\", 4);\n (\"great\", 5);\n ]\n in\n let tokenizer =\n word_level ~vocab ~unk_token:\"[UNK]\" ~pre:(Pre_tokenizer.whitespace ()) ()\n in\n\n let text = \"hello world is great\" in\n Printf.printf \"Text: %S\\n\" text;\n Printf.printf \"Length: %d tokens\\n\\n\"\n (Encoding.length (encode tokenizer text));\n print_encoding (encode tokenizer text);\n\n (* Show what happens with unknown words *)\n Printf.printf \"\\n--- Unknown words ---\\n\\n\";\n let text2 = \"hello universe\" in\n Printf.printf \"Text: %S\\n\" text2;\n Printf.printf \"Length: %d tokens\\n\\n\"\n (Encoding.length (encode tokenizer text2));\n print_encoding (encode tokenizer text2);\n\n (* WordPiece: subword tokens have word_ids linking to the source word *)\n Printf.printf \"\\n--- Subword tokens (WordPiece) ---\\n\\n\";\n let wp_vocab =\n [\n (\"[UNK]\", 0);\n (\"play\", 1);\n (\"##ing\", 2);\n (\"##ed\", 3);\n (\"un\", 4);\n (\"##happy\", 5);\n ]\n in\n let wp = wordpiece ~vocab:wp_vocab ~unk_token:\"[UNK]\" () in\n let text3 = \"playing\" in\n Printf.printf \"Text: %S\\n\" text3;\n Printf.printf \"Length: %d tokens\\n\\n\" (Encoding.length (encode wp text3));\n print_encoding (encode wp text3)\n"
},
{
"path": "packages/brot/examples/03-normalizers/README.md",
"content": "# `03-normalizers`\n\nText normalization before tokenization. Normalizers clean and standardize text\nso that surface variations (case, accents, whitespace) don't prevent vocabulary\nmatches.\n\n```bash\ndune exec brot/examples/03-normalizers/main.exe\n```\n\n## What You'll Learn\n\n- Unicode normalization: `nfc`, `nfkc`\n- Text transforms: `lowercase`, `strip_accents`, `strip`, `replace`, `prepend`\n- Model-specific normalization: `bert`\n- Composing normalizers with `sequence`\n- Applying normalizers directly with `Normalizer.apply`\n- How normalization affects tokenization results\n\n## Key Functions\n\n| Function | Purpose |\n| -------------------------- | ---------------------------------- |\n| `Normalizer.nfc` / `nfkc` | Unicode normalization forms |\n| `Normalizer.lowercase` | Unicode case folding |\n| `Normalizer.strip_accents` | Remove combining marks |\n| `Normalizer.strip` | Strip boundary whitespace |\n| `Normalizer.replace` | Regex-based replacement |\n| `Normalizer.prepend` | Prepend a string to non-empty text |\n| `Normalizer.bert` | BERT-specific normalizer |\n| `Normalizer.sequence` | Compose normalizers left-to-right |\n| `Normalizer.apply` | Apply a normalizer to a string |\n\n## Why Normalize?\n\nWithout normalization, `\"Hello\"`, `\"hello\"`, and `\"HELLO\"` are three different\ntokens. Normalization maps them all to `\"hello\"` so a single vocabulary entry\ncovers all cases. Similarly, `\"caf\\u{00E9}\"` and `\"cafe\"` can be unified by\nstripping accents.\n\n## Try It\n\n1. Add `Normalizer.nfkd` and see how it differs from `nfd`.\n2. Create a normalizer that replaces email addresses with ``.\n3. Try the BERT normalizer with Chinese characters.\n\n## Next Steps\n\nContinue to [04-pre-tokenizers](../04-pre-tokenizers/) to learn how text is\nsplit into fragments before vocabulary lookup.\n"
},
{
"path": "packages/brot/examples/03-normalizers/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/03-normalizers/main.ml",
"content": "(* Text normalization.\n\n Normalizers transform text before tokenization: lowercasing, accent removal,\n Unicode normalization, whitespace cleanup, and model-specific preprocessing.\n They are the first stage in the tokenization pipeline. *)\n\nopen Brot\n\nlet show name norm text =\n let result = Normalizer.apply norm text in\n Printf.printf \" %-20s %S -> %S\\n\" name text result\n\nlet () =\n Printf.printf \"=== Unicode Normalization ===\\n\\n\";\n show \"nfc\" Normalizer.nfc \"caf\\xc3\\xa9\";\n show \"nfkc\" Normalizer.nfkc \"\\xef\\xac\\x81\";\n\n (* fi ligature -> fi *)\n Printf.printf \"\\n=== Text Transforms ===\\n\\n\";\n show \"lowercase\" Normalizer.lowercase \"Hello WORLD\";\n show \"strip_accents\" Normalizer.strip_accents\n \"caf\\xc3\\xa9 r\\xc3\\xa9sum\\xc3\\xa9\";\n show \"strip\" (Normalizer.strip ()) \" hello \";\n show \"replace\"\n (Normalizer.replace ~pattern:\"\\\\d+\" ~replacement:\"\")\n \"I have 42 apples and 3 oranges\";\n show \"prepend\" (Normalizer.prepend \">> \") \"hello\";\n\n Printf.printf \"\\n=== Model-specific ===\\n\\n\";\n show \"bert (default)\" (Normalizer.bert ()) \"Hello WORLD!\";\n show \"bert (no lower)\" (Normalizer.bert ~lowercase:false ()) \"Hello WORLD!\";\n\n Printf.printf \"\\n=== Composition ===\\n\\n\";\n let composed =\n Normalizer.sequence\n [ Normalizer.nfd; Normalizer.strip_accents; Normalizer.lowercase ]\n in\n show \"nfd+strip+lower\" composed \"Caf\\xc3\\xa9 R\\xc3\\xa9sum\\xc3\\xa9\";\n show \"nfd+strip+lower\" composed \"HELLO\";\n\n Printf.printf \"\\n=== Effect on Tokenization ===\\n\\n\";\n let vocab =\n [ (\"hello\", 0); (\"world\", 1); (\"cafe\", 2); (\"resume\", 3); (\"\", 4) ]\n in\n let no_norm =\n word_level ~vocab ~unk_token:\"\" ~pre:(Pre_tokenizer.whitespace ()) ()\n in\n let with_norm =\n word_level ~vocab ~unk_token:\"\"\n ~pre:(Pre_tokenizer.whitespace ())\n ~normalizer:composed ()\n in\n\n let text = \"HELLO Caf\\xc3\\xa9\" in\n let enc1 = encode no_norm text in\n let enc2 = encode with_norm text in\n Printf.printf \" Text: %S\\n\" text;\n Printf.printf \" Without normalizer: [%s]\\n\"\n (String.concat \"; \"\n (List.map\n (fun s -> Printf.sprintf \"%S\" s)\n (Array.to_list (Encoding.tokens enc1))));\n Printf.printf \" With normalizer: [%s]\\n\"\n (String.concat \"; \"\n (List.map\n (fun s -> Printf.sprintf \"%S\" s)\n (Array.to_list (Encoding.tokens enc2))))\n"
},
{
"path": "packages/brot/examples/04-pre-tokenizers/README.md",
"content": "# `04-pre-tokenizers`\n\nPre-tokenization: splitting text into fragments before vocabulary lookup. Each\nfragment carries byte offsets into the original text.\n\n```bash\ndune exec brot/examples/04-pre-tokenizers/main.exe\n```\n\n## What You'll Learn\n\n- Common pre-tokenizers: `whitespace`, `whitespace_split`, `bert`\n- Punctuation and digit handling\n- Delimiter-based splitting: `char_delimiter`, `split`, `fixed_length`\n- SentencePiece-style `metaspace`\n- Composing pre-tokenizers with `sequence`\n- Using `Pre_tokenizer.pre_tokenize` to see fragments and offsets\n\n## Key Functions\n\n| Function | Purpose |\n| -------------------------------- | ------------------------------------------ |\n| `Pre_tokenizer.whitespace` | Pattern-based: `\\w+` and `[^\\w\\s]+` groups |\n| `Pre_tokenizer.whitespace_split` | Simple whitespace splitting |\n| `Pre_tokenizer.bert` | BERT-style: whitespace + punctuation + CJK |\n| `Pre_tokenizer.punctuation` | Isolate punctuation from words |\n| `Pre_tokenizer.digits` | Split on digit boundaries |\n| `Pre_tokenizer.char_delimiter` | Split on a single character |\n| `Pre_tokenizer.split` | Split on a literal string pattern |\n| `Pre_tokenizer.fixed_length` | Fixed-length character chunks |\n| `Pre_tokenizer.metaspace` | Replace spaces with visible markers |\n| `Pre_tokenizer.sequence` | Chain pre-tokenizers left-to-right |\n| `Pre_tokenizer.pre_tokenize` | Apply and get `(fragment, offsets)` list |\n\n## Pre-tokenizer vs Tokenizer\n\nPre-tokenization happens *before* the vocabulary-based algorithm (BPE,\nWordPiece, etc.). It determines the boundaries within which subword splitting\noperates. For example, with whitespace pre-tokenization, BPE will never merge\ntokens across word boundaries.\n\n## Try It\n\n1. Try `unicode_scripts` on text mixing Latin and CJK characters.\n2. Change the punctuation `behavior` to `` `Merged_with_previous `` or\n `` `Removed ``.\n3. Create a pre-tokenizer that splits on hyphens.\n\n## Next Steps\n\nContinue to [05-algorithms](../05-algorithms/) to see how different\ntokenization algorithms split the same text.\n"
},
{
"path": "packages/brot/examples/04-pre-tokenizers/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/04-pre-tokenizers/main.ml",
"content": "(* Pre-tokenization.\n\n Pre-tokenizers split text into fragments before vocabulary-based\n tokenization. Each fragment carries byte offsets into the original text.\n Different strategies produce different splits, affecting how subword\n algorithms see the input. *)\n\nopen Brot\n\nlet show name pre text =\n let result = Pre_tokenizer.pre_tokenize pre text in\n Printf.printf \" %-24s %S\\n\" name text;\n List.iter\n (fun (fragment, (s, e)) -> Printf.printf \" %S (%d, %d)\\n\" fragment s e)\n result;\n print_newline ()\n\nlet () =\n let text = \"Hello, world! It's 2026.\" in\n Printf.printf \"=== Common Pre-tokenizers ===\\n\\n\";\n Printf.printf \"Text: %S\\n\\n\" text;\n\n show \"whitespace\" (Pre_tokenizer.whitespace ()) text;\n show \"whitespace_split\" (Pre_tokenizer.whitespace_split ()) text;\n show \"bert\" (Pre_tokenizer.bert ()) text;\n show \"punctuation\" (Pre_tokenizer.punctuation ()) text;\n show \"digits (individual)\"\n (Pre_tokenizer.digits ~individual_digits:true ())\n text;\n show \"digits (grouped)\"\n (Pre_tokenizer.digits ~individual_digits:false ())\n text;\n\n Printf.printf \"=== Delimiter-based ===\\n\\n\";\n\n show \"char_delimiter ','\" (Pre_tokenizer.char_delimiter ',') \"a,b,c\";\n show \"split on '::'\" (Pre_tokenizer.split ~pattern:\"::\" ()) \"mod::func::arg\";\n show \"fixed_length 3\" (Pre_tokenizer.fixed_length 3) \"abcdefgh\";\n show \"metaspace\" (Pre_tokenizer.metaspace ()) \"Hello world today\";\n\n Printf.printf \"=== Composition ===\\n\\n\";\n\n let composed =\n Pre_tokenizer.sequence\n [\n Pre_tokenizer.whitespace_split ();\n Pre_tokenizer.punctuation ~behavior:`Isolated ();\n ]\n in\n show \"whitespace + punctuation\" composed text\n"
},
{
"path": "packages/brot/examples/05-algorithms/README.md",
"content": "# `05-algorithms`\n\nFive tokenization algorithms compared side-by-side. Each algorithm splits text\ndifferently based on its strategy.\n\n```bash\ndune exec brot/examples/05-algorithms/main.exe\n```\n\n## What You'll Learn\n\n- **BPE** (Byte Pair Encoding): merge-based subwords (GPT-2, RoBERTa)\n- **WordPiece**: greedy longest-match with `##` prefix (BERT)\n- **Unigram**: probabilistic segmentation (T5, mBART)\n- **Word-level**: one token per word, no subword splitting\n- **Character-level**: one token per byte, no vocabulary needed\n\n## Key Functions\n\n| Function | Purpose |\n| ----------------- | -------------------------------------- |\n| `Brot.bpe` | BPE tokenizer from vocab + merge rules |\n| `Brot.wordpiece` | WordPiece tokenizer from vocab |\n| `Brot.unigram` | Unigram tokenizer from vocab + scores |\n| `Brot.word_level` | Word-level tokenizer from vocab |\n| `Brot.chars` | Character-level tokenizer (no vocab) |\n| `Brot.vocab_size` | Number of vocabulary entries |\n\n## Algorithm Comparison\n\n| Algorithm | Subwords? | Unknown handling | Vocabulary |\n| ---------- | ------------------- | ------------------------ | ----------------------- |\n| BPE | Yes (merges) | Falls back to characters | `(string * int) list` |\n| WordPiece | Yes (`##` prefix) | `[UNK]` token | `(string * int) list` |\n| Unigram | Yes (probabilistic) | Lowest-score fallback | `(string * float) list` |\n| Word-level | No | `` token | `(string * int) list` |\n| Chars | No | N/A (all bytes valid) | None needed |\n\n## Try It\n\n1. Add more merge rules to the BPE tokenizer and see how it affects splitting.\n2. Try encoding a word not in the WordPiece vocabulary.\n3. Change the Unigram scores and observe how probabilities affect splitting.\n\n## Next Steps\n\nContinue to [06-special-tokens](../06-special-tokens/) to learn about special\ntokens and post-processing.\n"
},
{
"path": "packages/brot/examples/05-algorithms/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/05-algorithms/main.ml",
"content": "(* Tokenization algorithms.\n\n Five algorithms compared side-by-side: BPE (merge-based), WordPiece (greedy\n longest-match), Unigram (probabilistic), word-level (whole words), and\n character-level (per-byte). Each splits text differently. *)\n\nopen Brot\n\nlet show name tokenizer text =\n let encoding = encode tokenizer text in\n let tokens = Encoding.tokens encoding in\n let ids = Encoding.ids encoding in\n Printf.printf \" %-12s tokens=[%s] ids=[%s]\\n\" name\n (String.concat \", \"\n (List.map (fun s -> Printf.sprintf \"%S\" s) (Array.to_list tokens)))\n (String.concat \", \" (Array.to_list (Array.map string_of_int ids)))\n\nlet () =\n (* --- BPE: iterative merge-based subwords --- *)\n let bpe_tok =\n bpe\n ~vocab:\n [\n (\"p\", 0);\n (\"l\", 1);\n (\"a\", 2);\n (\"y\", 3);\n (\"i\", 4);\n (\"n\", 5);\n (\"g\", 6);\n (\"pl\", 7);\n (\"ay\", 8);\n (\"in\", 9);\n (\"ng\", 10);\n (\"play\", 11);\n (\"ing\", 12);\n (\"playing\", 13);\n ]\n ~merges:\n [\n (\"p\", \"l\");\n (\"a\", \"y\");\n (\"i\", \"n\");\n (\"n\", \"g\");\n (\"pl\", \"ay\");\n (\"in\", \"g\");\n (\"play\", \"ing\");\n ]\n ()\n in\n\n (* --- WordPiece: greedy longest-match with ## prefix --- *)\n let wp_tok =\n wordpiece\n ~vocab:\n [\n (\"[UNK]\", 0);\n (\"play\", 1);\n (\"##ing\", 2);\n (\"##ed\", 3);\n (\"run\", 4);\n (\"##ning\", 5);\n (\"un\", 6);\n (\"##known\", 7);\n ]\n ~unk_token:\"[UNK]\" ()\n in\n\n (* --- Unigram: probabilistic segmentation --- *)\n let uni_tok =\n unigram\n ~vocab:\n [\n (\"playing\", -0.5);\n (\"play\", -1.0);\n (\"ing\", -1.5);\n (\"p\", -3.0);\n (\"l\", -3.0);\n (\"a\", -3.0);\n (\"y\", -3.0);\n (\"i\", -3.0);\n (\"n\", -3.0);\n (\"g\", -3.0);\n ]\n ()\n in\n\n (* --- Word-level: whole words only --- *)\n let wl_tok =\n word_level\n ~vocab:[ (\"playing\", 0); (\"hello\", 1); (\"\", 2) ]\n ~unk_token:\"\"\n ~pre:(Pre_tokenizer.whitespace ())\n ()\n in\n\n (* --- Character-level: one byte per token --- *)\n let char_tok = chars () in\n\n Printf.printf \"=== Encoding %S ===\\n\\n\" \"playing\";\n show \"BPE\" bpe_tok \"playing\";\n show \"WordPiece\" wp_tok \"playing\";\n show \"Unigram\" uni_tok \"playing\";\n show \"Word-level\" wl_tok \"playing\";\n show \"Chars\" char_tok \"playing\";\n\n Printf.printf \"\\n=== Encoding %S ===\\n\\n\" \"running\";\n show \"WordPiece\" wp_tok \"running\";\n show \"Chars\" char_tok \"running\";\n\n Printf.printf \"\\n=== Encoding %S (unknown word) ===\\n\\n\" \"unknown\";\n show \"WordPiece\" wp_tok \"unknown\";\n show \"Word-level\" wl_tok \"unknown\";\n show \"Chars\" char_tok \"unknown\";\n\n Printf.printf \"\\n=== Vocabulary sizes ===\\n\\n\";\n Printf.printf \" BPE: %d\\n\" (vocab_size bpe_tok);\n Printf.printf \" WordPiece: %d\\n\" (vocab_size wp_tok);\n Printf.printf \" Unigram: %d\\n\" (vocab_size uni_tok);\n Printf.printf \" Word-level: %d\\n\" (vocab_size wl_tok);\n Printf.printf \" Chars: %d (byte range 0-255)\\n\" (vocab_size char_tok)\n"
},
{
"path": "packages/brot/examples/06-special-tokens/README.md",
"content": "# `06-special-tokens`\n\nSpecial tokens and post-processing. Post-processors insert tokens like `[CLS]`\nand `[SEP]` after tokenization, and assign type IDs for sentence-pair tasks.\n\n```bash\ndune exec brot/examples/06-special-tokens/main.exe\n```\n\n## What You'll Learn\n\n- Defining special tokens with `Brot.special`\n- BERT-style post-processing: `[CLS] A [SEP]` and `[CLS] A [SEP] B [SEP]`\n- Sentence-pair encoding with `encode ~pair`\n- Type IDs: 0 for first sequence, 1 for second\n- Template-based post-processing for custom formats\n- Skipping special tokens with `~add_special_tokens:false`\n\n## Key Functions\n\n| Function | Purpose |\n| ------------------------------ | ------------------------------------------- |\n| `Brot.special` | Define a special token configuration |\n| `Post_processor.bert` | BERT-style `[CLS] A [SEP] B [SEP]` |\n| `Post_processor.template` | Template-based with `$A`, `$B` placeholders |\n| `Brot.encode ~pair` | Encode a sentence pair |\n| `Encoding.type_ids` | Segment type IDs (0 or 1) |\n| `Encoding.special_tokens_mask` | 1 for special tokens, 0 for content |\n\n## BERT Post-processing\n\nFor a single sentence: `[CLS] tokens [SEP]`\nFor a sentence pair: `[CLS] A_tokens [SEP] B_tokens [SEP]`\n\nType IDs distinguish the two sequences:\n- First sequence (including `[CLS]` and first `[SEP]`): type_id = 0\n- Second sequence (including final `[SEP]`): type_id = 1\n\n## Try It\n\n1. Try the `roberta` post-processor with `` and ` ` tokens.\n2. Create a custom template with different special tokens.\n3. Encode a pair and check that `type_ids` correctly separates the segments.\n\n## Next Steps\n\nContinue to [07-padding-truncation](../07-padding-truncation/) to learn about\npreparing batches with uniform sequence lengths.\n"
},
{
"path": "packages/brot/examples/06-special-tokens/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/06-special-tokens/main.ml",
"content": "(* Special tokens and post-processing.\n\n Special tokens like [CLS] and [SEP] are inserted by post-processors after\n tokenization. They mark sequence boundaries and provide structure for model\n input. Sentence-pair encoding assigns different type IDs to each sequence. *)\n\nopen Brot\n\nlet print_encoding enc =\n let ids = Encoding.ids enc in\n let tokens = Encoding.tokens enc in\n let type_ids = Encoding.type_ids enc in\n let special = Encoding.special_tokens_mask enc in\n Printf.printf \" %-8s %-4s %-8s %-8s\\n\" \"Token\" \"ID\" \"Type_ID\" \"Special\";\n Printf.printf \" %s\\n\" (String.make 32 '-');\n for i = 0 to Encoding.length enc - 1 do\n Printf.printf \" %-8s %-4d %-8d %-8d\\n\" tokens.(i) ids.(i) type_ids.(i)\n special.(i)\n done\n\nlet () =\n let vocab =\n [\n (\"[UNK]\", 0);\n (\"[CLS]\", 1);\n (\"[SEP]\", 2);\n (\"hello\", 3);\n (\"world\", 4);\n (\"how\", 5);\n (\"are\", 6);\n (\"you\", 7);\n ]\n in\n let specials = List.map special [ \"[CLS]\"; \"[SEP]\"; \"[UNK]\" ] in\n let post = Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) () in\n let tokenizer =\n word_level ~vocab ~unk_token:\"[UNK]\" ~specials ~post\n ~pre:(Pre_tokenizer.whitespace ())\n ()\n in\n\n (* Single sentence: [CLS] A [SEP] *)\n Printf.printf \"=== Single Sentence ===\\n\";\n Printf.printf \"Text: \\\"hello world\\\"\\n\\n\";\n print_encoding (encode tokenizer \"hello world\");\n\n (* Sentence pair: [CLS] A [SEP] B [SEP] *)\n Printf.printf \"\\n=== Sentence Pair ===\\n\";\n Printf.printf \"A: \\\"hello world\\\", B: \\\"how are you\\\"\\n\\n\";\n print_encoding (encode tokenizer ~pair:\"how are you\" \"hello world\");\n\n (* Without special tokens *)\n Printf.printf \"\\n=== Without Special Tokens ===\\n\";\n Printf.printf \"Text: \\\"hello world\\\" (add_special_tokens=false)\\n\\n\";\n print_encoding (encode tokenizer ~add_special_tokens:false \"hello world\");\n\n (* Template-based post-processor *)\n Printf.printf \"\\n=== Template Post-processor ===\\n\";\n let template_post =\n Post_processor.template ~single:\"[CLS] $A [SEP]\"\n ~pair:\"[CLS] $A [SEP] $B:1 [SEP]:1\"\n ~special_tokens:[ (\"[CLS]\", 1); (\"[SEP]\", 2) ]\n ()\n in\n let tok2 =\n word_level ~vocab ~unk_token:\"[UNK]\" ~specials ~post:template_post\n ~pre:(Pre_tokenizer.whitespace ())\n ()\n in\n Printf.printf \"Template: \\\"[CLS] $A [SEP] $B:1 [SEP]:1\\\"\\n\";\n Printf.printf \"A: \\\"hello\\\", B: \\\"world\\\"\\n\\n\";\n print_encoding (encode tok2 ~pair:\"world\" \"hello\")\n"
},
{
"path": "packages/brot/examples/07-padding-truncation/README.md",
"content": "# `07-padding-truncation`\n\nPadding and truncation for batch processing. Models require uniform sequence\nlengths. Padding adds filler tokens; truncation trims long sequences.\n\n```bash\ndune exec brot/examples/07-padding-truncation/main.exe\n```\n\n## What You'll Learn\n\n- Fixed-length padding with `padding (`Fixed n)`\n- Batch-longest padding with `padding `Batch_longest`\n- Left vs right padding direction\n- Truncation with `truncation max_length`\n- Combining padding and truncation\n- Using `Encoding.attention_mask` to distinguish real tokens from padding\n\n## Key Functions\n\n| Function | Purpose |\n| ------------------------- | --------------------------------- |\n| `Brot.padding` | Create a padding configuration |\n| `Brot.truncation` | Create a truncation configuration |\n| `Brot.encode_batch` | Encode multiple texts at once |\n| `Encoding.attention_mask` | 1 for real tokens, 0 for padding |\n\n## Padding Strategies\n\n| Strategy | Behavior |\n| -------------------- | ------------------------------------------------------ |\n| `` `Fixed n `` | Every sequence padded to exactly `n` tokens |\n| `` `Batch_longest `` | All sequences padded to match the longest in the batch |\n| `` `To_multiple n `` | Pad to smallest multiple of `n` >= sequence length |\n\n## Try It\n\n1. Change the padding direction to `` `Left `` and observe where pad tokens appear.\n2. Try `padding (`To_multiple 4)` and see how lengths round up.\n3. Truncate from the left with `truncation ~direction:`Left 3`.\n\n## Next Steps\n\nContinue to [08-decoders](../08-decoders/) to learn how tokens are converted\nback to text.\n"
},
{
"path": "packages/brot/examples/07-padding-truncation/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/07-padding-truncation/main.ml",
"content": "(* Padding and truncation.\n\n Batch processing requires uniform sequence lengths. Padding extends short\n sequences with pad tokens; truncation trims long ones. The attention mask\n distinguishes real tokens from padding. *)\n\nopen Brot\n\nlet print_batch label encodings =\n Printf.printf \"%s\\n\" label;\n List.iteri\n (fun i enc ->\n let ids = Encoding.ids enc in\n let attn = Encoding.attention_mask enc in\n Printf.printf \" [%d] ids=[%s] attn=[%s]\\n\" i\n (String.concat \", \" (Array.to_list (Array.map string_of_int ids)))\n (String.concat \", \" (Array.to_list (Array.map string_of_int attn))))\n encodings;\n print_newline ()\n\nlet () =\n let vocab =\n [\n (\"[PAD]\", 0);\n (\"\", 1);\n (\"hello\", 2);\n (\"world\", 3);\n (\"how\", 4);\n (\"are\", 5);\n (\"you\", 6);\n (\"doing\", 7);\n (\"today\", 8);\n ]\n in\n let tokenizer =\n word_level ~vocab ~unk_token:\"\"\n ~specials:[ special \"[PAD]\" ]\n ~pad_token:\"[PAD]\"\n ~pre:(Pre_tokenizer.whitespace ())\n ()\n in\n\n let texts = [ \"hello\"; \"hello world\"; \"how are you doing today\" ] in\n Printf.printf \"Texts:\\n\";\n List.iteri (fun i t -> Printf.printf \" [%d] %S\\n\" i t) texts;\n print_newline ();\n\n (* No padding *)\n print_batch \"=== No Padding ===\" (encode_batch tokenizer texts);\n\n (* Fixed-length padding *)\n print_batch \"=== Fixed Padding (length=6) ===\"\n (encode_batch tokenizer ~padding:(padding (`Fixed 6)) texts);\n\n (* Batch-longest padding *)\n print_batch \"=== Batch Longest Padding ===\"\n (encode_batch tokenizer ~padding:(padding `Batch_longest) texts);\n\n (* Left padding *)\n print_batch \"=== Left Padding (length=6) ===\"\n (encode_batch tokenizer\n ~padding:(padding ~direction:`Left (`Fixed 6))\n texts);\n\n (* Truncation *)\n print_batch \"=== Truncation (max_length=3) ===\"\n (encode_batch tokenizer ~truncation:(truncation 3) texts);\n\n (* Padding + Truncation *)\n print_batch \"=== Padding + Truncation (pad=4, trunc=4) ===\"\n (encode_batch tokenizer\n ~padding:(padding (`Fixed 4))\n ~truncation:(truncation 4) texts)\n"
},
{
"path": "packages/brot/examples/08-decoders/README.md",
"content": "# `08-decoders`\n\nDecoders convert token strings back to natural text. Different tokenization\nschemes require different decoding strategies to produce clean output.\n\n```bash\ndune exec brot/examples/08-decoders/main.exe\n```\n\n## What You'll Learn\n\n- Per-token decoders: `wordpiece`, `bpe`, `metaspace`, `byte_fallback`\n- Collapsing decoders: `fuse`, `replace`\n- Composing decoders with `sequence`\n- Integrating a decoder with a tokenizer\n- Skipping special tokens during decoding\n\n## Key Functions\n\n| Function | Purpose |\n| ----------------------- | ------------------------------------ |\n| `Decoder.wordpiece` | Strip `##` prefix, join subwords |\n| `Decoder.bpe` | Strip word-end suffix, insert spaces |\n| `Decoder.metaspace` | Convert markers back to spaces |\n| `Decoder.byte_fallback` | Convert `<0xFF>` back to bytes |\n| `Decoder.fuse` | Concatenate all tokens |\n| `Decoder.replace` | String replacement |\n| `Decoder.sequence` | Chain decoders |\n| `Decoder.decode` | Apply decoder to token list |\n| `Brot.decode` | Full decode through tokenizer |\n\n## Per-token vs Collapsing\n\nSome decoders transform each token independently (per-token: `bpe`,\n`metaspace`, `byte_fallback`), while others combine the entire token list into\na single result (collapsing: `wordpiece`, `fuse`, `replace`). This matters\nwhen composing with `sequence`.\n\n## Try It\n\n1. Try `Decoder.ctc` for speech recognition CTC output.\n2. Compose `byte_fallback` with `fuse` and decode byte tokens.\n3. Use `Decoder.strip` to remove leading/trailing characters.\n\n## Next Steps\n\nContinue to [09-training](../09-training/) to learn how to train tokenizers\nfrom scratch.\n"
},
{
"path": "packages/brot/examples/08-decoders/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/08-decoders/main.ml",
"content": "(* Decoders.\n\n Decoders convert token strings back to natural text by reversing\n encoding-specific transformations: prefix/suffix removal, space insertion,\n byte-level decoding, and marker replacement. *)\n\nopen Brot\n\nlet show name decoder tokens =\n let result = Decoder.decode decoder tokens in\n Printf.printf \" %-22s [%s] -> %S\\n\" name\n (String.concat \"; \" (List.map (fun s -> Printf.sprintf \"%S\" s) tokens))\n result\n\nlet () =\n Printf.printf \"=== Per-token Decoders ===\\n\\n\";\n\n show \"wordpiece\" (Decoder.wordpiece ()) [ \"play\"; \"##ing\"; \"un\"; \"##happy\" ];\n\n show \"bpe (suffix=)\"\n (Decoder.bpe ~suffix:\"\" ())\n [ \"hel\"; \"lo\"; \"wor\"; \"ld\" ];\n\n show \"metaspace\" (Decoder.metaspace ())\n [ \"\\xe2\\x96\\x81Hello\"; \"\\xe2\\x96\\x81world\" ];\n\n show \"byte_fallback\" (Decoder.byte_fallback ()) [ \"hello\"; \"<0x21>\" ];\n\n Printf.printf \"\\n=== Collapsing Decoders ===\\n\\n\";\n\n show \"fuse\" (Decoder.fuse ()) [ \"h\"; \"e\"; \"l\"; \"l\"; \"o\" ];\n\n show \"replace ('_' -> ' ')\"\n (Decoder.replace ~pattern:\"_\" ~by:\" \" ())\n [ \"hello_world\" ];\n\n Printf.printf \"\\n=== Composed Decoder ===\\n\\n\";\n\n let composed =\n Decoder.sequence\n [ Decoder.wordpiece (); Decoder.replace ~pattern:\" \" ~by:\" \" () ]\n in\n show \"wordpiece + replace\" composed [ \"play\"; \"##ing\"; \"is\"; \"great\" ];\n\n Printf.printf \"\\n=== Integrated with Tokenizer ===\\n\\n\";\n\n let vocab =\n [\n (\"[UNK]\", 0);\n (\"[CLS]\", 1);\n (\"[SEP]\", 2);\n (\"play\", 3);\n (\"##ing\", 4);\n (\"##ed\", 5);\n (\"great\", 6);\n ]\n in\n let tokenizer =\n wordpiece ~vocab ~unk_token:\"[UNK]\"\n ~specials:[ special \"[CLS]\"; special \"[SEP]\" ]\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 1) ~sep:(\"[SEP]\", 2) ())\n ~decoder:(Decoder.wordpiece ()) ()\n in\n\n let text = \"playing\" in\n let encoding = encode tokenizer text in\n let ids = Encoding.ids encoding in\n Printf.printf \" Text: %S\\n\" text;\n Printf.printf \" Tokens: [%s]\\n\"\n (String.concat \"; \"\n (List.map\n (fun s -> Printf.sprintf \"%S\" s)\n (Array.to_list (Encoding.tokens encoding))));\n Printf.printf \" IDs: [%s]\\n\"\n (String.concat \"; \" (Array.to_list (Array.map string_of_int ids)));\n Printf.printf \" Decoded: %S\\n\" (decode tokenizer ids);\n Printf.printf \" Decoded (skip specials): %S\\n\"\n (decode tokenizer ~skip_special_tokens:true ids)\n"
},
{
"path": "packages/brot/examples/09-training/README.md",
"content": "# `09-training`\n\nTraining tokenizers from scratch. Given a text corpus, each algorithm learns a\nvocabulary tailored to the data.\n\n```bash\ndune exec brot/examples/09-training/main.exe\n```\n\n## What You'll Learn\n\n- Training BPE, WordPiece, word-level, and Unigram tokenizers\n- Controlling vocabulary size with `~vocab_size`\n- Adding special tokens during training\n- Inspecting the learned vocabulary\n\n## Key Functions\n\n| Function | Purpose |\n| ---------------------- | ------------------------------------------------ |\n| `Brot.train_bpe` | Train a BPE tokenizer (learns merge rules) |\n| `Brot.train_wordpiece` | Train a WordPiece tokenizer (learns subwords) |\n| `Brot.train_wordlevel` | Train a word-level tokenizer (collects words) |\n| `Brot.train_unigram` | Train a Unigram tokenizer (learns probabilities) |\n| `Brot.vocab_size` | Check learned vocabulary size |\n| `Brot.token_to_id` | Look up a token's ID |\n\n## Training Data\n\nTraining data is provided as `` `Seq (List.to_seq texts) `` for in-memory text\nor `` `Files [\"path1\"; \"path2\"] `` for files (one sentence per line).\n\n## Try It\n\n1. Add more sentences to the corpus and see how the vocabulary changes.\n2. Train with a smaller `~vocab_size` and observe more subword splitting.\n3. Use `~min_frequency:2` to exclude rare words.\n\n## Next Steps\n\nContinue to [10-bert-pipeline](../10-bert-pipeline/) to assemble a complete\nBERT-style tokenizer pipeline.\n"
},
{
"path": "packages/brot/examples/09-training/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/09-training/main.ml",
"content": "(* Training tokenizers.\n\n Train new tokenizers from a text corpus. Each algorithm learns a different\n vocabulary: BPE learns merge rules, WordPiece learns subword prefixes,\n word-level collects unique words, and Unigram learns token probabilities. *)\n\nopen Brot\n\nlet corpus =\n [\n \"the cat sat on the mat\";\n \"the dog sat on the log\";\n \"the cat and the dog are friends\";\n \"cats and dogs play together\";\n \"the cat plays with the dog\";\n \"playing in the park is fun\";\n \"the park has many cats and dogs\";\n \"friends play in the park together\";\n ]\n\nlet show_trained name tokenizer test_texts =\n Printf.printf \"--- %s (vocab_size=%d) ---\\n\" name (vocab_size tokenizer);\n List.iter\n (fun text ->\n let enc = encode tokenizer text in\n Printf.printf \" %S -> [%s]\\n\" text\n (String.concat \", \"\n (List.map\n (fun s -> Printf.sprintf \"%S\" s)\n (Array.to_list (Encoding.tokens enc)))))\n test_texts;\n print_newline ()\n\nlet () =\n let data = `Seq (List.to_seq corpus) in\n let test_texts = [ \"the cat plays\"; \"dogs are friends\" ] in\n\n Printf.printf \"Training corpus: %d sentences\\n\\n\" (List.length corpus);\n\n (* Train BPE: learns merge rules by iteratively combining frequent pairs *)\n let bpe_tok =\n train_bpe data ~vocab_size:100 ~show_progress:false\n ~pre:(Pre_tokenizer.whitespace ())\n in\n show_trained \"BPE\" bpe_tok test_texts;\n\n (* Train WordPiece: learns subword prefixes (## for continuation tokens) *)\n let wp_tok =\n train_wordpiece data ~vocab_size:100 ~show_progress:false\n ~pre:(Pre_tokenizer.whitespace ())\n in\n show_trained \"WordPiece\" wp_tok test_texts;\n\n (* Train word-level: each unique word is a token *)\n let wl_tok =\n train_wordlevel data ~vocab_size:50 ~show_progress:false\n ~pre:(Pre_tokenizer.whitespace ())\n in\n show_trained \"Word-level\" wl_tok test_texts;\n\n (* Train Unigram: probabilistic subword segmentation *)\n let uni_tok = train_unigram data ~vocab_size:100 ~show_progress:false in\n show_trained \"Unigram\" uni_tok test_texts;\n\n (* Training with special tokens *)\n Printf.printf \"=== Training with Special Tokens ===\\n\\n\";\n let wp_with_specials =\n train_wordpiece data ~vocab_size:100 ~show_progress:false\n ~pre:(Pre_tokenizer.whitespace ())\n ~specials:[ special \"[CLS]\"; special \"[SEP]\"; special \"[PAD]\" ]\n ~pad_token:\"[PAD]\"\n in\n Printf.printf \"WordPiece with specials (vocab=%d):\\n\"\n (vocab_size wp_with_specials);\n let show_id tok name =\n Printf.printf \" %s id = %s\\n\" name\n (match token_to_id tok name with\n | Some id -> string_of_int id\n | None -> \"N/A\")\n in\n show_id wp_with_specials \"[CLS]\";\n show_id wp_with_specials \"[SEP]\";\n show_id wp_with_specials \"[PAD]\";\n\n (* Add a post-processor to insert special tokens during encoding *)\n Printf.printf \"\\n Encoding with post-processor:\\n\";\n let wp_full =\n train_wordpiece data ~vocab_size:100 ~show_progress:false\n ~pre:(Pre_tokenizer.whitespace ())\n ~post:\n (Post_processor.bert\n ~cls:(\"[CLS]\", Option.get (token_to_id wp_with_specials \"[CLS]\"))\n ~sep:(\"[SEP]\", Option.get (token_to_id wp_with_specials \"[SEP]\"))\n ())\n ~specials:[ special \"[CLS]\"; special \"[SEP]\"; special \"[PAD]\" ]\n ~pad_token:\"[PAD]\"\n in\n let enc = encode wp_full \"the cat plays\" in\n Printf.printf \" %S -> [%s]\\n\" \"the cat plays\"\n (String.concat \", \"\n (List.map\n (fun s -> Printf.sprintf \"%S\" s)\n (Array.to_list (Encoding.tokens enc))))\n"
},
{
"path": "packages/brot/examples/10-bert-pipeline/README.md",
"content": "# `10-bert-pipeline`\n\nComplete BERT-style tokenizer pipeline. Assembles all stages: normalizer,\npre-tokenizer, WordPiece algorithm, post-processor, decoder, special tokens,\npadding, and truncation.\n\n```bash\ndune exec brot/examples/10-bert-pipeline/main.exe\n```\n\n## What You'll Learn\n\n- Assembling a full tokenization pipeline\n- How all stages work together end-to-end\n- Single sentence and sentence-pair encoding\n- Batch encoding with padding\n- Sentence-pair batch encoding with `encode_pairs_batch`\n- Decoding with and without special tokens\n- Inspecting tokenizer configuration with `Brot.pp`\n\n## Key Functions\n\n| Function | Purpose |\n| ---------------------------------- | --------------------------------------------- |\n| `Brot.wordpiece` | Full pipeline constructor |\n| `Normalizer.bert` | BERT normalizer (lowercase, clean, CJK) |\n| `Pre_tokenizer.bert` | BERT pre-tokenizer (whitespace + punctuation) |\n| `Post_processor.bert` | Insert `[CLS]` and `[SEP]` tokens |\n| `Decoder.wordpiece` | Reverse `##` prefix joining |\n| `Brot.encode ~pair` | Encode a sentence pair |\n| `Brot.encode_pairs_batch` | Batch-encode sentence pairs |\n| `Brot.decode ~skip_special_tokens` | Decode without `[CLS]`/`[SEP]` |\n| `Brot.pp` | Pretty-print tokenizer configuration |\n\n## The Full Pipeline\n\n```\nInput text\n |\n v\nNormalizer.bert -- lowercase, clean control chars, pad CJK\n |\n v\nPre_tokenizer.bert -- split on whitespace, isolate punctuation\n |\n v\nWordPiece model -- greedy longest-match subword splitting\n |\n v\nPost_processor.bert -- insert [CLS] and [SEP], set type_ids\n |\n v\nEncoding.t -- ids, tokens, offsets, type_ids, attention_mask\n```\n\n## Try It\n\n1. Encode text with accented characters and see the normalizer at work.\n2. Change `Post_processor.bert` to `Post_processor.roberta` with `` and\n ` ` tokens for a RoBERTa-style pipeline.\n3. Use `save_pretrained` to export the tokenizer and reload it with\n `from_file`.\n\n## Further Reading\n\n- [gpt2_tokenizer](../x-gpt2-tokenizer/) -- loading a real GPT-2 tokenizer\n from HuggingFace model files\n"
},
{
"path": "packages/brot/examples/10-bert-pipeline/dune",
"content": "(executable\n (name main)\n (libraries brot))\n"
},
{
"path": "packages/brot/examples/10-bert-pipeline/main.ml",
"content": "(* BERT-style pipeline.\n\n Assembles all pipeline stages into a complete BERT-style tokenizer:\n normalizer, pre-tokenizer, WordPiece algorithm, post-processor, decoder,\n special tokens, padding, and truncation. *)\n\nopen Brot\n\nlet print_encoding label enc =\n let tokens = Encoding.tokens enc in\n let ids = Encoding.ids enc in\n let type_ids = Encoding.type_ids enc in\n let attn = Encoding.attention_mask enc in\n Printf.printf \"%s\\n\" label;\n Printf.printf \" tokens: [%s]\\n\"\n (String.concat \", \"\n (List.map (fun s -> Printf.sprintf \"%S\" s) (Array.to_list tokens)));\n Printf.printf \" ids: [%s]\\n\"\n (String.concat \", \" (Array.to_list (Array.map string_of_int ids)));\n Printf.printf \" type_ids: [%s]\\n\"\n (String.concat \", \" (Array.to_list (Array.map string_of_int type_ids)));\n Printf.printf \" attn_mask: [%s]\\n\"\n (String.concat \", \" (Array.to_list (Array.map string_of_int attn)));\n print_newline ()\n\nlet () =\n (* Build a BERT-style vocabulary *)\n let vocab =\n [\n (\"[PAD]\", 0);\n (\"[UNK]\", 1);\n (\"[CLS]\", 2);\n (\"[SEP]\", 3);\n (\"the\", 4);\n (\"cat\", 5);\n (\"sat\", 6);\n (\"on\", 7);\n (\"mat\", 8);\n (\"dog\", 9);\n (\"play\", 10);\n (\"##ing\", 11);\n (\"##ed\", 12);\n (\"is\", 13);\n (\"a\", 14);\n (\"good\", 15);\n (\"great\", 16);\n (\"un\", 17);\n (\"##happy\", 18);\n (\"friend\", 19);\n (\"##s\", 20);\n (\"how\", 21);\n (\"are\", 22);\n (\"you\", 23);\n ]\n in\n let specials = List.map special [ \"[PAD]\"; \"[UNK]\"; \"[CLS]\"; \"[SEP]\" ] in\n\n (* Assemble the full pipeline *)\n let tokenizer =\n wordpiece ~vocab ~unk_token:\"[UNK]\"\n ~normalizer:(Normalizer.bert ~lowercase:true ())\n ~pre:(Pre_tokenizer.bert ())\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 2) ~sep:(\"[SEP]\", 3) ())\n ~decoder:(Decoder.wordpiece ()) ~specials ~pad_token:\"[PAD]\" ()\n in\n\n (* Inspect the tokenizer *)\n Printf.printf \"=== Tokenizer Configuration ===\\n\";\n Format.printf \"%a@.@.\" pp tokenizer;\n\n (* Single sentence *)\n Printf.printf \"=== Single Sentence ===\\n\\n\";\n print_encoding \"\\\"The Cat is Playing\\\"\"\n (encode tokenizer \"The Cat is Playing\");\n\n (* Sentence pair *)\n Printf.printf \"=== Sentence Pair ===\\n\\n\";\n print_encoding \"A: \\\"the cat sat\\\", B: \\\"how are you\\\"\"\n (encode tokenizer ~pair:\"how are you\" \"the cat sat\");\n\n (* Batch with padding *)\n Printf.printf \"=== Padded Batch ===\\n\\n\";\n let batch =\n encode_batch tokenizer ~padding:(padding `Batch_longest)\n [ \"the cat\"; \"the cat sat on a mat\"; \"good\" ]\n in\n List.iteri (fun i enc -> print_encoding (Printf.sprintf \"[%d]\" i) enc) batch;\n\n (* Sentence pairs batch with padding and truncation *)\n Printf.printf \"=== Sentence Pairs (pad=12, trunc=12) ===\\n\\n\";\n let pairs =\n encode_pairs_batch tokenizer\n ~padding:(padding (`Fixed 12))\n ~truncation:(truncation 12)\n [ (\"the cat sat\", \"how are you\"); (\"good dog\", \"is a friend\") ]\n in\n List.iteri\n (fun i enc -> print_encoding (Printf.sprintf \"pair[%d]\" i) enc)\n pairs;\n\n (* Decoding *)\n Printf.printf \"=== Decoding ===\\n\\n\";\n let enc = encode tokenizer ~pair:\"how are you\" \"the cat sat\" in\n let ids = Encoding.ids enc in\n Printf.printf \" Full decode: %S\\n\" (decode tokenizer ids);\n Printf.printf \" Skip specials: %S\\n\"\n (decode tokenizer ~skip_special_tokens:true ids)\n"
},
{
"path": "packages/brot/examples/README.md",
"content": "# Brot Examples\n\nLearn Brot through progressively complex examples. Start with `01-encode-decode`\nand work through the numbered examples in order.\n\n## Examples\n\n| Example | Concept | Key Functions |\n|---------|---------|---------------|\n| [`01-encode-decode`](./01-encode-decode/) | Text to IDs and back | `bpe`, `encode`, `decode` |\n| [`02-encoding-fields`](./02-encoding-fields/) | Encoding metadata | `Encoding.ids`, `.tokens`, `.offsets` |\n| [`03-normalizers`](./03-normalizers/) | Text normalization | `Normalizer.lowercase`, `.bert`, `.sequence` |\n| [`04-pre-tokenizers`](./04-pre-tokenizers/) | Splitting before vocab | `Pre_tokenizer.whitespace`, `.bert`, `.sequence` |\n| [`05-algorithms`](./05-algorithms/) | Algorithm comparison | `bpe`, `wordpiece`, `unigram`, `word_level`, `chars` |\n| [`06-special-tokens`](./06-special-tokens/) | Special tokens and post-processing | `Post_processor.bert`, `.template`, `encode ~pair` |\n| [`07-padding-truncation`](./07-padding-truncation/) | Batch preparation | `padding`, `truncation`, `encode_batch` |\n| [`08-decoders`](./08-decoders/) | Tokens back to text | `Decoder.wordpiece`, `.bpe`, `.fuse`, `.sequence` |\n| [`09-training`](./09-training/) | Train from scratch | `train_bpe`, `train_wordpiece`, `train_unigram` |\n| [`10-bert-pipeline`](./10-bert-pipeline/) | Full BERT pipeline | All stages assembled end-to-end |\n\nAdvanced:\n\n- [**x-gpt2-tokenizer**](./x-gpt2-tokenizer/): Loading a real GPT-2 tokenizer\n from HuggingFace model files\n\n## Running Examples\n\nAll examples can be run with:\n\n```bash\ndune exec brot/examples//main.exe\n```\n\nFor example:\n\n```bash\ndune exec brot/examples/01-encode-decode/main.exe\n```\n\n## Quick Reference\n\n### Encode and Decode\n\n```ocaml\nopen Brot\n\nlet tokenizer = bpe ~vocab:[(\"hello\", 0); ...] ~merges:[...] () in\nlet encoding = encode tokenizer \"hello world\" in\nlet ids = Encoding.ids encoding in\nlet text = decode tokenizer ids\n```\n\n### Full Pipeline\n\n```ocaml\nlet tokenizer =\n wordpiece ~vocab\n ~normalizer:(Normalizer.bert ~lowercase:true ())\n ~pre:(Pre_tokenizer.bert ())\n ~post:(Post_processor.bert ~cls:(\"[CLS]\", 2) ~sep:(\"[SEP]\", 3) ())\n ~decoder:(Decoder.wordpiece ())\n ~specials:(List.map special [ \"[CLS]\"; \"[SEP]\"; \"[PAD]\" ])\n ~pad_token:\"[PAD]\" ()\n```\n\n### Train from Text\n\n```ocaml\nlet tokenizer =\n train_bpe (`Seq (List.to_seq texts)) ~vocab_size:1000\n```\n"
},
{
"path": "packages/brot/examples/x-gpt2-tokenizer/README.md",
"content": "# `x-gpt2-tokenizer`\n\nLoading a real GPT-2 tokenizer from HuggingFace model files. This example\ndownloads GPT-2's vocabulary and merges, builds the full byte-level BPE\npipeline, and demonstrates encoding, decoding, and subword inspection.\n\n```bash\ndune exec brot/examples/x-gpt2-tokenizer/main.exe\n```\n\n## What You'll Learn\n\n- Loading a pre-trained tokenizer from vocabulary and merge files\n- Building a byte-level BPE pipeline with `from_model_file`\n- Encoding text and inspecting tokens, IDs, and offsets\n- Decoding token IDs back to text\n- Subword splitting on real vocabulary\n- Batch encoding multiple texts\n\n## Key Functions\n\n| Function | Purpose |\n| -------------------------- | ----------------------------------------------- |\n| `Brot.from_model_file` | Load tokenizer from vocab.json and merges.txt |\n| `Pre_tokenizer.byte_level` | GPT-2 style byte-level pre-tokenizer |\n| `Decoder.byte_level` | Corresponding byte-level decoder |\n| `Brot.encode` | Encode text to an `Encoding.t` |\n| `Brot.decode` | Decode token IDs back to text |\n| `Brot.encode_batch` | Encode multiple texts at once |\n| `Encoding.tokens` | Token strings from an encoding |\n| `Encoding.ids` | Token IDs from an encoding |\n| `Encoding.offsets` | Byte offset pairs mapping tokens to source text |\n\n## Prerequisites\n\nThis example downloads GPT-2 model files from HuggingFace on first run\n(~1 MB total). Files are cached in `/tmp/brot_gpt2/`.\n\n## Output Walkthrough\n\n```\nVocabulary: 50257 tokens\n\nText: \"Hello world! GPT-2 is amazing.\"\nTokens: [\"Hello\"; \" world\"; \"!\"; \" GPT\"; \"-\"; \"2\"; \" is\"; \" amazing\"; \".\"]\nIDs: [15496; 995; 0; 402; 12; 17; 318; 4998; 13]\nDecoded: \"Hello world! GPT-2 is amazing.\"\nRound-trip: true\n\n=== Subword Splitting ===\n\n \"tokenization\" -> 3 tokens: [\"token\", \"ization\"]\n \"transformer\" -> 1 tokens: [\"transformer\"]\n ...\n\n=== Batch Encoding ===\n\n \"The quick brown fox\" -> 4 tokens\n \"jumps over the lazy dog\" -> 5 tokens\n \"Machine learning is fun\" -> 4 tokens\n\n=== Token Offsets ===\n\n Text: \"Hello, world!\"\n Hello offsets=(0, 5) source=\"Hello\"\n , offsets=(5, 6) source=\",\"\n ...\n```\n\n## Try It\n\n1. Change the input text and see how GPT-2 tokenizes different sentences.\n2. Try words with unusual spellings to see subword splitting in action.\n3. Compare the token count for English text vs other languages.\n\n## See Also\n\n- [01-encode-decode](../01-encode-decode/) for basic encoding and decoding\n- [05-algorithms](../05-algorithms/) for comparing tokenization algorithms\n- [08-decoders](../08-decoders/) for decoder options\n"
},
{
"path": "packages/brot/examples/x-gpt2-tokenizer/dune",
"content": "(executable\n (name main)\n (libraries brot nx unix))\n"
},
{
"path": "packages/brot/examples/x-gpt2-tokenizer/main.ml",
"content": "(* Loading a real GPT-2 tokenizer.\n\n Downloads GPT-2's vocabulary and merge files from HuggingFace, builds the\n full byte-level BPE pipeline, and demonstrates encoding, decoding, and\n subword inspection on real-world text. *)\n\nopen Brot\n\nlet download url dest =\n if not (Sys.file_exists dest) then (\n Printf.printf \"Downloading %s...\\n%!\" (Filename.basename dest);\n let cmd =\n Printf.sprintf \"curl -L --fail -s -o %s %s\" (Filename.quote dest)\n (Filename.quote url)\n in\n match Unix.system cmd with\n | Unix.WEXITED 0 -> ()\n | _ -> failwith (Printf.sprintf \"Failed to download %s\" url))\n\nlet () =\n (* Download GPT-2 model files *)\n let cache = \"/tmp/brot_gpt2\" in\n if not (Sys.file_exists cache) then Sys.mkdir cache 0o755;\n let vocab_file = Filename.concat cache \"vocab.json\" in\n let merges_file = Filename.concat cache \"merges.txt\" in\n download \"https://huggingface.co/gpt2/raw/main/vocab.json\" vocab_file;\n download \"https://huggingface.co/gpt2/raw/main/merges.txt\" merges_file;\n\n (* Build the GPT-2 tokenizer: BPE with byte-level pre-tokenizer *)\n let tokenizer =\n from_model_file ~vocab:vocab_file ~merges:merges_file\n ~pre:(Pre_tokenizer.byte_level ~add_prefix_space:false ())\n ~decoder:(Decoder.byte_level ()) ()\n in\n Printf.printf \"\\nVocabulary: %d tokens\\n\\n\" (vocab_size tokenizer);\n\n (* Encode text *)\n let text = \"Hello world! GPT-2 is amazing.\" in\n let enc = encode tokenizer text in\n Printf.printf \"Text: %S\\n\" text;\n Printf.printf \"Tokens: [%s]\\n\"\n (String.concat \"; \"\n (List.map\n (fun s -> Printf.sprintf \"%S\" s)\n (Array.to_list (Encoding.tokens enc))));\n Printf.printf \"IDs: [%s]\\n\"\n (String.concat \"; \"\n (Array.to_list (Array.map string_of_int (Encoding.ids enc))));\n\n (* Decode back *)\n let decoded = decode tokenizer (Encoding.ids enc) in\n Printf.printf \"Decoded: %S\\n\" decoded;\n Printf.printf \"Round-trip: %b\\n\\n\" (String.equal text decoded);\n\n (* Subword splitting: see how a long word is broken down *)\n Printf.printf \"=== Subword Splitting ===\\n\\n\";\n List.iter\n (fun word ->\n let e = encode tokenizer word in\n let tokens = Encoding.tokens e in\n Printf.printf \" %-20s -> %d tokens: [%s]\\n\" (Printf.sprintf \"%S\" word)\n (Array.length tokens)\n (String.concat \", \"\n (List.map (fun s -> Printf.sprintf \"%S\" s) (Array.to_list tokens))))\n [ \"tokenization\"; \"transformer\"; \"GPT\"; \"Hello\"; \"supercalifragilistic\" ];\n\n (* Batch encoding *)\n Printf.printf \"\\n=== Batch Encoding ===\\n\\n\";\n let texts =\n [\n \"The quick brown fox\";\n \"jumps over the lazy dog\";\n \"Machine learning is fun\";\n ]\n in\n let batch = encode_batch tokenizer texts in\n List.iter2\n (fun text enc ->\n Printf.printf \" %-30s -> %d tokens\\n\" (Printf.sprintf \"%S\" text)\n (Encoding.length enc))\n texts batch;\n\n (* Offsets: map tokens back to source text *)\n Printf.printf \"\\n=== Token Offsets ===\\n\\n\";\n let text2 = \"Hello, world!\" in\n let enc2 = encode tokenizer text2 in\n Printf.printf \"Text: %S\\n\" text2;\n let tokens = Encoding.tokens enc2 in\n let offsets = Encoding.offsets enc2 in\n for i = 0 to Encoding.length enc2 - 1 do\n let s, e = offsets.(i) in\n Printf.printf \" %-8s offsets=(%d, %d) source=%S\\n\" tokens.(i) s e\n (String.sub text2 s (e - s))\n done\n"
},
{
"path": "packages/brot/lib/bpe.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet list_drop n l =\n let rec aux i = function\n | _ :: l when i < n -> aux (i + 1) l\n | rest -> rest\n in\n if n <= 0 then l else aux 0 l\n\ntype vocab = (string, int) Hashtbl.t\ntype merges = (string * string) list\n\n(* Open-addressing hash table for merge lookups. Returns int directly (no option\n allocation). -1 = not found. *)\nmodule Merge_map = struct\n type t = { keys : int array; values : int array; mask : int }\n\n let[@inline] hash key =\n let h = key * 0x1B873593 in\n h lxor (h lsr 16)\n\n let create entries =\n let n = List.length entries in\n let cap = ref 16 in\n while !cap < n * 4 do\n cap := !cap * 2\n done;\n let mask = !cap - 1 in\n let keys = Array.make !cap (-1) in\n let values = Array.make !cap 0 in\n List.iter\n (fun (key, value) ->\n let h = ref (hash key land mask) in\n while Array.unsafe_get keys !h >= 0 do\n h := (!h + 1) land mask\n done;\n Array.unsafe_set keys !h key;\n Array.unsafe_set values !h value)\n entries;\n { keys; values; mask }\n\n let[@inline] find t key =\n let mask = t.mask in\n let keys = t.keys in\n let h = ref (hash key land mask) in\n let k = ref (Array.unsafe_get keys !h) in\n while !k <> key && !k >= 0 do\n h := (!h + 1) land mask;\n k := Array.unsafe_get keys !h\n done;\n if !k = key then Array.unsafe_get t.values !h else -1\n\n let fold f t acc =\n let keys = t.keys in\n let values = t.values in\n let len = Array.length keys in\n let acc = ref acc in\n for i = 0 to len - 1 do\n let k = Array.unsafe_get keys i in\n if k >= 0 then acc := f k (Array.unsafe_get values i) !acc\n done;\n !acc\nend\n\nlet[@inline] merge_key a b = (a lsl 21) lor b\nlet[@inline] pack_merge rank new_id = (rank lsl 21) lor new_id\nlet[@inline] merge_rank v = v lsr 21\nlet[@inline] merge_new_id v = v land 0x1FFFFF\n\ntype word = {\n sym_c : int array;\n sym_prev : int array;\n sym_next : int array;\n sym_len : int array;\n mutable size : int;\n}\n\n(* Specialized min-heap for BPE merges using parallel arrays (no tuple allocation).\n Ordered by (rank, position) — lower rank first, then lower position. *)\n(* Min-heap with packed comparison key: (rank lsl 21) lor pos.\n Single int comparison for sift operations, 2 arrays instead of 3. *)\nmodule Merge_queue = struct\n type t = {\n mutable keys : int array;\n mutable new_ids : int array;\n mutable size : int;\n mutable pop_key : int;\n mutable pop_new_id : int;\n mutable skip_keys : int array;\n mutable skip_new_ids : int array;\n mutable skip_size : int;\n }\n\n let create cap =\n let cap = max 16 cap in\n {\n keys = Array.make cap 0;\n new_ids = Array.make cap 0;\n size = 0;\n pop_key = 0;\n pop_new_id = 0;\n skip_keys = [||];\n skip_new_ids = [||];\n skip_size = 0;\n }\n\n let[@inline] pack_key rank pos = (rank lsl 21) lor pos\n\n let sift_up t idx =\n let keys = t.keys in\n let new_ids = t.new_ids in\n let key = Array.unsafe_get keys idx in\n let nid = Array.unsafe_get new_ids idx in\n let i = ref idx in\n let cont = ref (!i > 0) in\n while !cont do\n let p = (!i - 1) asr 1 in\n if key < Array.unsafe_get keys p then (\n Array.unsafe_set keys !i (Array.unsafe_get keys p);\n Array.unsafe_set new_ids !i (Array.unsafe_get new_ids p);\n i := p;\n cont := !i > 0)\n else cont := false\n done;\n Array.unsafe_set keys !i key;\n Array.unsafe_set new_ids !i nid\n\n let sift_down t idx =\n let keys = t.keys in\n let new_ids = t.new_ids in\n let size = t.size in\n let key = Array.unsafe_get keys idx in\n let nid = Array.unsafe_get new_ids idx in\n let i = ref idx in\n let continue_ = ref true in\n while !continue_ do\n let l = (2 * !i) + 1 in\n if l >= size then continue_ := false\n else begin\n let r = l + 1 in\n let smallest =\n if r < size && Array.unsafe_get keys r < Array.unsafe_get keys l then\n r\n else l\n in\n if Array.unsafe_get keys smallest < key then (\n Array.unsafe_set keys !i (Array.unsafe_get keys smallest);\n Array.unsafe_set new_ids !i (Array.unsafe_get new_ids smallest);\n i := smallest)\n else continue_ := false\n end\n done;\n Array.unsafe_set keys !i key;\n Array.unsafe_set new_ids !i nid\n\n let push t rank pos new_id =\n let s = t.size in\n if s = Array.length t.keys then begin\n let new_cap = max 16 (s * 2) in\n let grow a =\n let b = Array.make new_cap 0 in\n Array.blit a 0 b 0 s;\n b\n in\n t.keys <- grow t.keys;\n t.new_ids <- grow t.new_ids\n end;\n Array.unsafe_set t.keys s (pack_key rank pos);\n Array.unsafe_set t.new_ids s new_id;\n t.size <- s + 1;\n sift_up t s\n\n let pop t =\n if t.size = 0 then false\n else begin\n t.pop_key <- Array.unsafe_get t.keys 0;\n t.pop_new_id <- Array.unsafe_get t.new_ids 0;\n t.size <- t.size - 1;\n if t.size > 0 then begin\n Array.unsafe_set t.keys 0 (Array.unsafe_get t.keys t.size);\n Array.unsafe_set t.new_ids 0 (Array.unsafe_get t.new_ids t.size);\n sift_down t 0\n end;\n true\n end\nend\n\ntype token = { id : int; value : string; offsets : int * int }\n\n(* Direct-mapped bounded cache: hash key to slot, newest entry wins. Fixed\n memory, no eviction logic, no unbounded growth. *)\ntype cache = {\n cache_keys : string array;\n cache_vals : word array;\n cache_mask : int;\n}\n\nlet empty_word =\n { sym_c = [||]; sym_prev = [||]; sym_next = [||]; sym_len = [||]; size = 0 }\n\nlet create_cache capacity =\n (* Round up to power of 2 *)\n let cap = ref 16 in\n while !cap < capacity do\n cap := !cap * 2\n done;\n {\n cache_keys = Array.make !cap \"\";\n cache_vals = Array.make !cap empty_word;\n cache_mask = !cap - 1;\n }\n\nlet[@inline] cache_find c key =\n let h = Hashtbl.hash key land c.cache_mask in\n if String.equal (Array.unsafe_get c.cache_keys h) key then\n Array.unsafe_get c.cache_vals h\n else empty_word\n\nlet[@inline] cache_add c key value =\n let h = Hashtbl.hash key land c.cache_mask in\n Array.unsafe_set c.cache_keys h key;\n Array.unsafe_set c.cache_vals h value\n\ntype t = {\n vocab : vocab;\n vocab_r : string array;\n merges : Merge_map.t;\n cache : cache option;\n dropout : float option;\n unk_token : string option;\n continuing_subword_prefix : string option;\n end_of_word_suffix : string option;\n fuse_unk : bool;\n byte_fallback : bool;\n ignore_merges : bool;\n ascii_to_id : int array;\n byte_fallback_ids : int array;\n char_to_id : Merge_map.t;\n prefixed_ascii_to_id : int array;\n prefixed_char_to_id : Merge_map.t;\n unk_id : int;\n mutable work_word : word;\n mutable work_queue : Merge_queue.t;\n work_in_use : bool Atomic.t;\n}\n\nlet create_word capacity =\n let cap = max 16 capacity in\n {\n sym_c = Array.make cap 0;\n sym_prev = Array.make cap 0;\n sym_next = Array.make cap 0;\n sym_len = Array.make cap 0;\n size = 0;\n }\n\nlet ensure_word_capacity word capacity =\n if Array.length word.sym_c >= capacity then begin\n word.size <- 0;\n word\n end\n else create_word capacity\n\nlet ensure_queue_capacity queue capacity =\n let cap = max 16 capacity in\n if Array.length queue.Merge_queue.keys >= cap then begin\n queue.Merge_queue.size <- 0;\n queue\n end\n else Merge_queue.create cap\n\nlet[@inline] add_symbol word c byte_len =\n let s = word.size in\n let prev = if s > 0 then s - 1 else -1 in\n Array.unsafe_set word.sym_c s c;\n Array.unsafe_set word.sym_prev s prev;\n Array.unsafe_set word.sym_next s (-1);\n Array.unsafe_set word.sym_len s byte_len;\n if prev >= 0 then Array.unsafe_set word.sym_next prev s;\n word.size <- s + 1\n\nlet apply_merges model dropout word queue =\n let p = match dropout with Some p -> p | None -> 0.0 in\n let use_dropout = p > 0.0 in\n let merges = model.merges in\n let sym_c = word.sym_c in\n let sym_prev = word.sym_prev in\n let sym_next = word.sym_next in\n let sym_len = word.sym_len in\n for i = 0 to word.size - 2 do\n let key =\n merge_key (Array.unsafe_get sym_c i) (Array.unsafe_get sym_c (i + 1))\n in\n let packed = Merge_map.find merges key in\n if packed >= 0 then\n Merge_queue.push queue (merge_rank packed) i (merge_new_id packed)\n done;\n queue.skip_size <- 0;\n while Merge_queue.pop queue do\n let pkey = queue.pop_key in\n let pos = pkey land 0x1FFFFF in\n let new_id = queue.pop_new_id in\n if Array.unsafe_get sym_len pos > 0 then begin\n let next_pos = Array.unsafe_get sym_next pos in\n if next_pos >= 0 then begin\n let key =\n merge_key\n (Array.unsafe_get sym_c pos)\n (Array.unsafe_get sym_c next_pos)\n in\n let packed = Merge_map.find merges key in\n if packed >= 0 && merge_new_id packed = new_id then\n begin if use_dropout && Random.float 1.0 < p then begin\n let s = queue.skip_size in\n if s = Array.length queue.skip_keys then begin\n let new_cap = max 8 (s * 2) in\n let grow old =\n let a = Array.make new_cap 0 in\n if s > 0 then Array.blit old 0 a 0 s;\n a\n in\n queue.skip_keys <- grow queue.skip_keys;\n queue.skip_new_ids <- grow queue.skip_new_ids\n end;\n Array.unsafe_set queue.skip_keys s pkey;\n Array.unsafe_set queue.skip_new_ids s new_id;\n queue.skip_size <- s + 1\n end\n else begin\n for i = 0 to queue.skip_size - 1 do\n Merge_queue.push queue\n (Array.unsafe_get queue.skip_keys i lsr 21)\n (Array.unsafe_get queue.skip_keys i land 0x1FFFFF)\n (Array.unsafe_get queue.skip_new_ids i)\n done;\n queue.skip_size <- 0;\n Array.unsafe_set sym_c pos new_id;\n Array.unsafe_set sym_len pos\n (Array.unsafe_get sym_len pos + Array.unsafe_get sym_len next_pos);\n Array.unsafe_set sym_next pos (Array.unsafe_get sym_next next_pos);\n Array.unsafe_set sym_len next_pos 0;\n let new_next = Array.unsafe_get sym_next pos in\n if new_next >= 0 then Array.unsafe_set sym_prev new_next pos;\n let prev = Array.unsafe_get sym_prev pos in\n if prev >= 0 then begin\n let k =\n merge_key\n (Array.unsafe_get sym_c prev)\n (Array.unsafe_get sym_c pos)\n in\n let v = Merge_map.find merges k in\n if v >= 0 then\n Merge_queue.push queue (merge_rank v) prev (merge_new_id v)\n end;\n let next = Array.unsafe_get sym_next pos in\n if next >= 0 then begin\n let k =\n merge_key\n (Array.unsafe_get sym_c pos)\n (Array.unsafe_get sym_c next)\n in\n let v = Merge_map.find merges k in\n if v >= 0 then\n Merge_queue.push queue (merge_rank v) pos (merge_new_id v)\n end\n end\n end\n end\n end\n done;\n (* Compact using linked-list traversal: O(N_final) instead of O(N_original) *)\n let j = ref 0 in\n let cur = ref 0 in\n while !cur >= 0 do\n if !j <> !cur then begin\n Array.unsafe_set sym_c !j (Array.unsafe_get sym_c !cur);\n Array.unsafe_set sym_len !j (Array.unsafe_get sym_len !cur)\n end;\n incr j;\n cur := Array.unsafe_get sym_next !cur\n done;\n word.size <- !j\n\nlet utf8_byte_len_table =\n Array.init 256 (fun b ->\n if b land 0x80 = 0 then 1\n else if b land 0xE0 = 0xC0 then 2\n else if b land 0xF0 = 0xE0 then 3\n else if b land 0xF8 = 0xF0 then 4\n else 1)\n\nlet[@inline] utf8_byte_len b = Array.unsafe_get utf8_byte_len_table b\n\nlet[@inline] pack_char_key text pos byte_len =\n let b0 = Char.code (String.unsafe_get text pos) in\n match byte_len with\n | 1 -> b0\n | 2 -> (b0 lsl 8) lor Char.code (String.unsafe_get text (pos + 1))\n | 3 ->\n (b0 lsl 16)\n lor (Char.code (String.unsafe_get text (pos + 1)) lsl 8)\n lor Char.code (String.unsafe_get text (pos + 2))\n | _ ->\n (b0 lsl 24)\n lor (Char.code (String.unsafe_get text (pos + 1)) lsl 16)\n lor (Char.code (String.unsafe_get text (pos + 2)) lsl 8)\n lor Char.code (String.unsafe_get text (pos + 3))\n\n(* Try emitting byte fallback tokens for [byte_len] bytes starting at [src]\n offset [offset]. Returns true if all bytes had fallback IDs. *)\nlet try_byte_fallback model word flush_unk src offset byte_len =\n let all_found = ref true in\n for i = 0 to byte_len - 1 do\n if\n Array.unsafe_get model.byte_fallback_ids\n (Char.code (String.unsafe_get src (offset + i)))\n < 0\n then all_found := false\n done;\n if !all_found then begin\n flush_unk ();\n for i = 0 to byte_len - 1 do\n add_symbol word\n (Array.unsafe_get model.byte_fallback_ids\n (Char.code (String.unsafe_get src (offset + i))))\n 1\n done;\n true\n end\n else false\n\n(* No prefix/suffix — avoids all per-character string allocation for ASCII via\n pre-computed lookup tables. *)\nlet init_word_fast model word text text_len =\n let pos = ref 0 in\n let pending_unk_id = ref (-1) in\n let pending_unk_len = ref 0 in\n let flush_unk () =\n if !pending_unk_id >= 0 then begin\n add_symbol word !pending_unk_id !pending_unk_len;\n pending_unk_id := -1;\n pending_unk_len := 0\n end\n in\n let handle_unk byte_len =\n if model.unk_id >= 0 then\n begin if model.fuse_unk then\n begin if !pending_unk_id >= 0 then\n pending_unk_len := !pending_unk_len + byte_len\n else begin\n pending_unk_id := model.unk_id;\n pending_unk_len := byte_len\n end\n end\n else begin\n flush_unk ();\n add_symbol word model.unk_id byte_len\n end\n end\n in\n while !pos < text_len do\n let b = Char.code (String.unsafe_get text !pos) in\n if b < 128 then begin\n let id = Array.unsafe_get model.ascii_to_id b in\n if id >= 0 then begin\n flush_unk ();\n add_symbol word id 1\n end\n else if model.byte_fallback then begin\n let fbid = Array.unsafe_get model.byte_fallback_ids b in\n if fbid >= 0 then begin\n flush_unk ();\n add_symbol word fbid 1\n end\n else handle_unk 1\n end\n else handle_unk 1;\n incr pos\n end\n else begin\n let byte_len = utf8_byte_len b in\n let key = pack_char_key text !pos byte_len in\n let id = Merge_map.find model.char_to_id key in\n if id >= 0 then begin\n flush_unk ();\n add_symbol word id byte_len\n end\n else if model.byte_fallback then\n begin if not (try_byte_fallback model word flush_unk text !pos byte_len)\n then handle_unk byte_len\n end\n else handle_unk byte_len;\n pos := !pos + byte_len\n end\n done;\n flush_unk ()\n\n(* Models with continuing_subword_prefix or end_of_word_suffix *)\nlet init_word_slow model word text text_len =\n let pos = ref 0 in\n let pending_unk_id = ref (-1) in\n let pending_unk_len = ref 0 in\n let flush_unk () =\n if !pending_unk_id >= 0 then begin\n add_symbol word !pending_unk_id !pending_unk_len;\n pending_unk_id := -1;\n pending_unk_len := 0\n end\n in\n let handle_unk byte_len =\n if model.unk_id >= 0 then\n begin if model.fuse_unk then\n begin if !pending_unk_id >= 0 then\n pending_unk_len := !pending_unk_len + byte_len\n else begin\n pending_unk_id := model.unk_id;\n pending_unk_len := byte_len\n end\n end\n else begin\n flush_unk ();\n add_symbol word model.unk_id byte_len\n end\n end\n in\n let has_prefix = model.continuing_subword_prefix <> None in\n let has_suffix = model.end_of_word_suffix <> None in\n while !pos < text_len do\n let b = Char.code (String.unsafe_get text !pos) in\n let byte_len = utf8_byte_len b in\n if b land 0xC0 = 0x80 then pos := !pos + 1\n else begin\n let start = !pos in\n let is_first = start = 0 in\n let is_last = !pos + byte_len >= text_len in\n pos := !pos + byte_len;\n (* Suffix only applies at word boundaries (first-not-last or\n last-not-first), never to middle chars and never to single-char\n words *)\n let needs_string = has_suffix && is_first <> is_last in\n if needs_string then begin\n (* Slow path: suffix involved, at most 2x per word *)\n let char_str = String.sub text start byte_len in\n let token_str =\n match\n ( is_first,\n is_last,\n model.continuing_subword_prefix,\n model.end_of_word_suffix )\n with\n | true, false, _, Some suffix -> char_str ^ suffix\n | true, false, _, None -> char_str\n | false, true, Some prefix, Some suffix -> prefix ^ char_str ^ suffix\n | false, true, Some prefix, None -> prefix ^ char_str\n | false, true, None, Some suffix -> char_str ^ suffix\n | false, true, None, None -> char_str\n | _, _, _, _ -> char_str\n in\n match Hashtbl.find_opt model.vocab token_str with\n | Some id ->\n flush_unk ();\n add_symbol word id byte_len\n | None ->\n if model.byte_fallback then\n begin if\n not (try_byte_fallback model word flush_unk text start byte_len)\n then handle_unk byte_len\n end\n else handle_unk byte_len\n end\n else begin\n (* Fast path: no suffix, use packed-int lookup (zero allocation) *)\n let needs_prefix = has_prefix && not is_first in\n let id =\n if needs_prefix then\n if b < 128 then Array.unsafe_get model.prefixed_ascii_to_id b\n else\n Merge_map.find model.prefixed_char_to_id\n (pack_char_key text start byte_len)\n else if b < 128 then Array.unsafe_get model.ascii_to_id b\n else\n Merge_map.find model.char_to_id (pack_char_key text start byte_len)\n in\n if id >= 0 then begin\n flush_unk ();\n add_symbol word id byte_len\n end\n else if model.byte_fallback then\n begin if\n not (try_byte_fallback model word flush_unk text start byte_len)\n then handle_unk byte_len\n end\n else handle_unk byte_len\n end\n end\n done;\n flush_unk ()\n\nlet merge_word model text =\n let text_len = String.length text in\n let owned = Atomic.compare_and_set model.work_in_use false true in\n let word, queue =\n if owned then begin\n let w = ensure_word_capacity model.work_word text_len in\n model.work_word <- w;\n let q = ensure_queue_capacity model.work_queue text_len in\n model.work_queue <- q;\n (w, q)\n end\n else (create_word text_len, Merge_queue.create text_len)\n in\n if model.continuing_subword_prefix = None && model.end_of_word_suffix = None\n then init_word_fast model word text text_len\n else init_word_slow model word text text_len;\n apply_merges model model.dropout word queue;\n if owned then begin\n let n = word.size in\n let sym_c = Array.make n 0 in\n let sym_len = Array.make n 0 in\n Array.blit word.sym_c 0 sym_c 0 n;\n Array.blit word.sym_len 0 sym_len 0 n;\n Atomic.set model.work_in_use false;\n { sym_c; sym_prev = [||]; sym_next = [||]; sym_len; size = n }\n end\n else word\n\nlet word_to_tokens model word =\n let offset = ref 0 in\n List.init word.size (fun i ->\n let id = Array.unsafe_get word.sym_c i in\n let vr = model.vocab_r in\n let value =\n if id >= 0 && id < Array.length vr then Array.unsafe_get vr id\n else \"\"\n in\n let start = !offset in\n let end_ = start + Array.unsafe_get word.sym_len i in\n offset := end_;\n { id; value; offsets = (start, end_) })\n\nlet word_to_ids word =\n Array.init word.size (fun i -> Array.unsafe_get word.sym_c i)\n\nlet word_to_encoding model word ~type_id =\n let n = word.size in\n let ids = Array.make n 0 in\n let tokens = Array.make n \"\" in\n let offsets = Array.make n (0, 0) in\n let offset = ref 0 in\n for i = 0 to n - 1 do\n let id = Array.unsafe_get word.sym_c i in\n Array.unsafe_set ids i id;\n let vr = model.vocab_r in\n Array.unsafe_set tokens i\n (if id >= 0 && id < Array.length vr then Array.unsafe_get vr id\n else \"\");\n let start = !offset in\n let end_ = start + Array.unsafe_get word.sym_len i in\n Array.unsafe_set offsets i (start, end_);\n offset := end_\n done;\n Encoding.create ~ids ~type_ids:(Array.make n type_id) ~tokens\n ~words:(Array.make n None) ~offsets ~special_tokens_mask:(Array.make n 0)\n ~attention_mask:(Array.make n 1) ()\n\nlet get_word model text =\n if model.ignore_merges then merge_word model text\n else\n match model.cache with\n | Some cache when String.length text < 4096 ->\n let cached = cache_find cache text in\n if cached.size > 0 then cached\n else\n let word = merge_word model text in\n cache_add cache text word;\n word\n | _ -> merge_word model text\n\nlet tokenize model text =\n if String.length text = 0 then []\n else\n match Hashtbl.find_opt model.vocab text with\n | Some id -> [ { id; value = text; offsets = (0, String.length text) } ]\n | None -> word_to_tokens model (get_word model text)\n\nlet tokenize_ids model text =\n if String.length text = 0 then [||]\n else\n match Hashtbl.find_opt model.vocab text with\n | Some id -> [| id |]\n | None -> word_to_ids (get_word model text)\n\nlet tokenize_encoding model text ~type_id =\n if String.length text = 0 then Encoding.empty\n else\n match Hashtbl.find_opt model.vocab text with\n | Some id ->\n Encoding.token ~id ~token:text\n ~offset:(0, String.length text)\n ~type_id ~special:false\n | None -> word_to_encoding model (get_word model text) ~type_id\n\nlet token_to_id model token = Hashtbl.find_opt model.vocab token\n\nlet id_to_token model id =\n if id >= 0 && id < Array.length model.vocab_r then\n Some (Array.unsafe_get model.vocab_r id)\n else None\n\nlet get_vocab model = Hashtbl.fold (fun k v acc -> (k, v) :: acc) model.vocab []\nlet get_vocab_size model = Hashtbl.length model.vocab\nlet get_unk_token model = model.unk_token\nlet get_continuing_subword_prefix model = model.continuing_subword_prefix\nlet get_end_of_word_suffix model = model.end_of_word_suffix\n\nlet get_merges model =\n Merge_map.fold\n (fun key packed acc ->\n let a_id = key lsr 21 in\n let b_id = key land 0x1FFFFF in\n let rank = merge_rank packed in\n let vr = model.vocab_r in\n let vr_len = Array.length vr in\n if a_id >= 0 && a_id < vr_len && b_id >= 0 && b_id < vr_len then\n (rank, (Array.unsafe_get vr a_id, Array.unsafe_get vr b_id)) :: acc\n else acc)\n model.merges []\n |> List.sort (fun (r1, _) (r2, _) -> Int.compare r1 r2)\n |> List.map snd\n\nlet convert_merges_to_merge_map vocab merges continuing_subword_prefix =\n let csp_str =\n match continuing_subword_prefix with Some p -> p | None -> \"\"\n in\n let csp_len = String.length csp_str in\n List.mapi\n (fun rank (a, b) ->\n match (Hashtbl.find_opt vocab a, Hashtbl.find_opt vocab b) with\n | Some a_id, Some b_id -> (\n let alen = String.length a in\n let blen = String.length b in\n let new_token =\n if\n csp_len > 0 && blen > csp_len\n && String.starts_with ~prefix:csp_str b\n then (\n let brest = blen - csp_len in\n let s = Bytes.create (alen + brest) in\n Bytes.blit_string a 0 s 0 alen;\n Bytes.blit_string b csp_len s alen brest;\n Bytes.unsafe_to_string s)\n else\n let s = Bytes.create (alen + blen) in\n Bytes.blit_string a 0 s 0 alen;\n Bytes.blit_string b 0 s alen blen;\n Bytes.unsafe_to_string s\n in\n match Hashtbl.find_opt vocab new_token with\n | Some new_id -> Some ((a_id, b_id), pack_merge rank new_id)\n | None ->\n failwith\n (Printf.sprintf \"Merge token '%s' not in vocabulary\" new_token))\n | _ ->\n failwith\n (Printf.sprintf \"Merge tokens ('%s', '%s') not in vocabulary\" a b))\n merges\n |> List.filter_map Fun.id\n |> fun entries ->\n Merge_map.create\n (List.map\n (fun ((a_id, b_id), packed) -> (merge_key a_id b_id, packed))\n entries)\n\nlet create ~vocab ~merges ?(cache_capacity = 10000) ?dropout ?unk_token\n ?continuing_subword_prefix ?end_of_word_suffix ?(fuse_unk = false)\n ?(byte_fallback = false) ?(ignore_merges = false) () : t =\n let max_id = Hashtbl.fold (fun _ id acc -> max id acc) vocab (-1) in\n let vocab_r = Array.make (max_id + 1) \"\" in\n Hashtbl.iter (fun k v -> Array.unsafe_set vocab_r v k) vocab;\n let cache =\n if cache_capacity = 0 then None else Some (create_cache cache_capacity)\n in\n let merges =\n convert_merges_to_merge_map vocab merges continuing_subword_prefix\n in\n let ascii_to_id = Array.make 128 (-1) in\n for i = 0 to 127 do\n let s = String.make 1 (Char.chr i) in\n match Hashtbl.find_opt vocab s with\n | Some id -> ascii_to_id.(i) <- id\n | None -> ()\n done;\n let byte_fallback_ids = Array.make 256 (-1) in\n for i = 0 to 255 do\n let hex = Printf.sprintf \"<0x%02X>\" i in\n match Hashtbl.find_opt vocab hex with\n | Some id -> byte_fallback_ids.(i) <- id\n | None -> ()\n done;\n (* Build packed-int char lookup table for zero-allocation multi-byte lookup *)\n let char_entries = ref [] in\n Hashtbl.iter\n (fun key id ->\n let len = String.length key in\n if len >= 1 && len <= 4 then begin\n let b0 = Char.code (String.unsafe_get key 0) in\n let expected_len = utf8_byte_len b0 in\n if expected_len = len then\n let packed =\n match len with\n | 1 -> b0\n | 2 -> (b0 lsl 8) lor Char.code (String.unsafe_get key 1)\n | 3 ->\n (b0 lsl 16)\n lor (Char.code (String.unsafe_get key 1) lsl 8)\n lor Char.code (String.unsafe_get key 2)\n | _ ->\n (b0 lsl 24)\n lor (Char.code (String.unsafe_get key 1) lsl 16)\n lor (Char.code (String.unsafe_get key 2) lsl 8)\n lor Char.code (String.unsafe_get key 3)\n in\n char_entries := (packed, id) :: !char_entries\n end)\n vocab;\n let char_to_id = Merge_map.create !char_entries in\n (* Build prefixed char lookup tables for zero-allocation init_word_slow *)\n let prefixed_ascii_to_id = Array.make 128 (-1) in\n let prefixed_char_entries = ref [] in\n (match continuing_subword_prefix with\n | Some prefix ->\n for i = 0 to 127 do\n let s = prefix ^ String.make 1 (Char.chr i) in\n match Hashtbl.find_opt vocab s with\n | Some id -> prefixed_ascii_to_id.(i) <- id\n | None -> ()\n done;\n Hashtbl.iter\n (fun key id ->\n let plen = String.length prefix in\n let klen = String.length key in\n if klen > plen && String.sub key 0 plen = prefix then begin\n let rest_len = klen - plen in\n if rest_len >= 2 && rest_len <= 4 then begin\n let b0 = Char.code (String.unsafe_get key plen) in\n let expected = utf8_byte_len b0 in\n if expected = rest_len then\n let packed = pack_char_key key plen rest_len in\n prefixed_char_entries := (packed, id) :: !prefixed_char_entries\n end\n end)\n vocab\n | None -> ());\n let prefixed_char_to_id = Merge_map.create !prefixed_char_entries in\n let unk_id =\n match unk_token with\n | Some unk -> (\n match Hashtbl.find_opt vocab unk with Some id -> id | None -> -1)\n | None -> -1\n in\n {\n vocab;\n vocab_r;\n merges;\n cache;\n dropout;\n unk_token;\n continuing_subword_prefix;\n end_of_word_suffix;\n fuse_unk;\n byte_fallback;\n ignore_merges;\n ascii_to_id;\n byte_fallback_ids;\n char_to_id;\n prefixed_ascii_to_id;\n prefixed_char_to_id;\n unk_id;\n work_word = create_word 16;\n work_queue = Merge_queue.create 16;\n work_in_use = Atomic.make false;\n }\n\nlet json_of_string s =\n match Jsont_bytesrw.decode_string Jsont.json s with\n | Ok v -> v\n | Error e -> failwith e\n\nlet json_to_string j =\n match Jsont_bytesrw.encode_string ~format:Jsont.Minify Jsont.json j with\n | Ok s -> s\n | Error e -> failwith e\n\nlet read_files ~vocab_file ~merges_file =\n let vocab_json =\n let ic = open_in vocab_file in\n let content =\n Fun.protect\n ~finally:(fun () -> close_in ic)\n (fun () -> really_input_string ic (in_channel_length ic))\n in\n json_of_string content\n in\n let vocab = Hashtbl.create 1024 in\n (match vocab_json with\n | Jsont.Object (mems, _) ->\n List.iter\n (fun ((k, _), v) ->\n match v with\n | Jsont.Number (f, _) -> Hashtbl.add vocab k (int_of_float f)\n | _ -> failwith \"Invalid vocab format\")\n mems\n | _ -> failwith \"Invalid vocab.json format\");\n let merges =\n let ic = open_in merges_file in\n Fun.protect\n ~finally:(fun () -> close_in ic)\n (fun () ->\n let merges = ref [] in\n (try\n while true do\n let line = input_line ic in\n (* Skip empty lines and comment lines that start with #version *)\n if\n String.length line > 0\n && not (String.starts_with ~prefix:\"#version\" line)\n then\n match String.split_on_char ' ' line with\n | [ a; b ] -> merges := (a, b) :: !merges\n | _ -> failwith (Printf.sprintf \"Invalid merge line: %s\" line)\n done\n with End_of_file -> ());\n List.rev !merges)\n in\n (vocab, merges)\n\nlet from_files ~vocab_file ~merges_file =\n let vocab, merges = read_files ~vocab_file ~merges_file in\n create ~vocab ~merges ()\n\nlet save model ~path ?name () =\n let vocab_file =\n match name with\n | Some n -> Filename.concat path (Printf.sprintf \"%s-vocab.json\" n)\n | None -> Filename.concat path \"vocab.json\"\n in\n let merges_file =\n match name with\n | Some n -> Filename.concat path (Printf.sprintf \"%s-merges.txt\" n)\n | None -> Filename.concat path \"merges.txt\"\n in\n let vocab_items =\n Hashtbl.fold (fun k v acc -> (k, v) :: acc) model.vocab []\n |> List.sort (fun (_, a) (_, b) -> compare a b)\n in\n let vocab_json =\n Jsont.Json.object'\n (List.map\n (fun (k, v) -> (Jsont.Json.name k, Jsont.Json.int v))\n vocab_items)\n in\n let oc = open_out vocab_file in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc (json_to_string vocab_json));\n let oc = open_out merges_file in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () ->\n output_string oc \"#version: 0.2\\n\";\n let merges_list =\n Merge_map.fold\n (fun key packed acc ->\n let a_id = key lsr 21 in\n let b_id = key land 0x1FFFFF in\n let rank = merge_rank packed in\n let vr = model.vocab_r in\n let vr_len = Array.length vr in\n if a_id >= 0 && a_id < vr_len && b_id >= 0 && b_id < vr_len then\n (rank, Array.unsafe_get vr a_id, Array.unsafe_get vr b_id) :: acc\n else acc)\n model.merges []\n |> List.sort (fun (r1, _, _) (r2, _, _) -> compare r1 r2)\n in\n List.iter (fun (_, a, b) -> Printf.fprintf oc \"%s %s\\n\" a b) merges_list)\n\nlet train ~min_frequency ~vocab_size ~show_progress ~special_tokens\n ~limit_alphabet ~initial_alphabet ~continuing_subword_prefix\n ~end_of_word_suffix ~max_token_length texts existing =\n let _ = (show_progress, existing) in\n\n (* Count words from texts *)\n let word_counts = Hashtbl.create 10000 in\n List.iter\n (fun text ->\n let words = String.split_on_char ' ' text in\n List.iter\n (fun word ->\n if String.length word > 0 then\n Hashtbl.replace word_counts word\n (1 + try Hashtbl.find word_counts word with Not_found -> 0))\n words)\n texts;\n\n let compute_pair_counts words_copy =\n let pair_counts = Hashtbl.create 10000 in\n Hashtbl.iter\n (fun word count ->\n let chars = String.split_on_char ' ' word in\n for i = 0 to List.length chars - 2 do\n let a = List.nth chars i in\n let b = List.nth chars (i + 1) in\n let pair = (a, b) in\n Hashtbl.replace pair_counts pair\n (count + try Hashtbl.find pair_counts pair with Not_found -> 0)\n done)\n words_copy;\n pair_counts\n in\n\n (* Build vocabulary *)\n let vocab = Hashtbl.create 10000 in\n let vocab_size_ref = ref 0 in\n List.iter\n (fun token ->\n if not (Hashtbl.mem vocab token) then (\n Hashtbl.add vocab token !vocab_size_ref;\n incr vocab_size_ref))\n special_tokens;\n\n (* Build alphabet *)\n let alphabet = Hashtbl.create 10000 in\n Hashtbl.iter\n (fun word count ->\n let len = String.length word in\n let buf = Buffer.create 4 in\n let rec loop i =\n if i >= len then ()\n else\n let d = String.get_utf_8_uchar word i in\n let n = Uchar.utf_decode_length d in\n if Uchar.utf_decode_is_valid d then (\n let u = Uchar.utf_decode_uchar d in\n Buffer.clear buf;\n Buffer.add_utf_8_uchar buf u;\n let char_str = Buffer.contents buf in\n Hashtbl.replace alphabet char_str\n (count + try Hashtbl.find alphabet char_str with Not_found -> 0));\n loop (i + n)\n in\n loop 0)\n word_counts;\n\n List.iter\n (fun c ->\n let char_str = String.make 1 c in\n Hashtbl.replace alphabet char_str max_int)\n initial_alphabet;\n\n let kept = Hashtbl.fold (fun k v acc -> (k, v) :: acc) alphabet [] in\n let kept = List.sort (fun (_, v1) (_, v2) -> compare v1 v2) kept in\n let to_remove =\n match limit_alphabet with\n | Some limit -> max 0 (List.length kept - limit)\n | None -> 0\n in\n let kept = list_drop to_remove kept in\n let kept = List.sort (fun (k1, _) (k2, _) -> compare k1 k2) kept in\n let csp_str =\n match continuing_subword_prefix with Some p -> p | None -> \"\"\n in\n let csp_len = String.length csp_str in\n List.iter\n (fun (c, _) ->\n if not (Hashtbl.mem vocab c) then (\n Hashtbl.add vocab c !vocab_size_ref;\n incr vocab_size_ref);\n if csp_len > 0 then (\n let clen = String.length c in\n let s = Bytes.create (csp_len + clen) in\n Bytes.blit_string csp_str 0 s 0 csp_len;\n Bytes.blit_string c 0 s csp_len clen;\n let prefixed = Bytes.unsafe_to_string s in\n if not (Hashtbl.mem vocab prefixed) then (\n Hashtbl.add vocab prefixed !vocab_size_ref;\n incr vocab_size_ref)))\n kept;\n\n (* Learn merges *)\n let merges = ref [] in\n let words_copy = ref (Hashtbl.create (Hashtbl.length word_counts)) in\n Hashtbl.iter\n (fun word count ->\n let len = String.length word in\n let chars = ref [] in\n let buf = Buffer.create 8 in\n let is_first = ref true in\n let rec loop i =\n if i >= len then ()\n else\n let d = String.get_utf_8_uchar word i in\n let n = Uchar.utf_decode_length d in\n if Uchar.utf_decode_is_valid d then (\n let u = Uchar.utf_decode_uchar d in\n Buffer.clear buf;\n if csp_len > 0 && not !is_first then Buffer.add_string buf csp_str;\n Buffer.add_utf_8_uchar buf u;\n is_first := false;\n chars := Buffer.contents buf :: !chars);\n loop (i + n)\n in\n loop 0;\n let separated = String.concat \" \" (List.rev !chars) in\n Hashtbl.add !words_copy separated count)\n word_counts;\n\n while !vocab_size_ref < vocab_size do\n let pair_counts = compute_pair_counts !words_copy in\n let best_pair = ref None in\n let best_count = ref (-1) in\n let best_pair_tie = ref (\"\", \"\") in\n Hashtbl.iter\n (fun pair count ->\n if count > !best_count then (\n best_count := count;\n best_pair := Some pair;\n best_pair_tie := pair)\n else if count = !best_count then\n if compare pair !best_pair_tie < 0 then best_pair_tie := pair)\n pair_counts;\n match !best_pair with\n | None -> vocab_size_ref := vocab_size\n | Some (a, b) ->\n if !best_count < min_frequency then vocab_size_ref := vocab_size\n else\n let blen = String.length b in\n let new_token =\n if\n csp_len > 0 && blen > csp_len\n && String.starts_with ~prefix:csp_str b\n then (\n let alen = String.length a in\n let brest = blen - csp_len in\n let s = Bytes.create (alen + brest) in\n Bytes.blit_string a 0 s 0 alen;\n Bytes.blit_string b csp_len s alen brest;\n Bytes.unsafe_to_string s)\n else a ^ b\n in\n let skip =\n match max_token_length with\n | Some l when String.length new_token > l -> true\n | _ -> false\n in\n if not skip then (\n if not (Hashtbl.mem vocab new_token) then (\n Hashtbl.add vocab new_token !vocab_size_ref;\n incr vocab_size_ref);\n merges := (a, b) :: !merges;\n let new_words = Hashtbl.create (Hashtbl.length !words_copy) in\n let pat = a ^ \" \" ^ b in\n let pat_len = String.length pat in\n Hashtbl.iter\n (fun word count ->\n let wlen = String.length word in\n if wlen < pat_len then Hashtbl.add new_words word count\n else\n let buf = Buffer.create wlen in\n let pos = ref 0 in\n let changed = ref false in\n while !pos <= wlen - pat_len do\n let at_boundary =\n (!pos = 0\n || Char.equal (String.unsafe_get word (!pos - 1)) ' ')\n && (!pos + pat_len = wlen\n || Char.equal\n (String.unsafe_get word (!pos + pat_len))\n ' ')\n in\n if at_boundary && String.sub word !pos pat_len = pat then (\n Buffer.add_string buf new_token;\n pos := !pos + pat_len;\n changed := true)\n else (\n Buffer.add_char buf (String.unsafe_get word !pos);\n incr pos)\n done;\n if !changed then (\n Buffer.add_substring buf word !pos (wlen - !pos);\n Hashtbl.add new_words (Buffer.contents buf) count)\n else Hashtbl.add new_words word count)\n !words_copy;\n words_copy := new_words)\n done;\n\n let trained_model =\n create ~vocab ~merges:(List.rev !merges) ?continuing_subword_prefix\n ?end_of_word_suffix ()\n in\n (trained_model, special_tokens)\n"
},
{
"path": "packages/brot/lib/bpe.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** BPE (Byte Pair Encoding) tokenization model.\n\n {b Internal module.} Iteratively merges the most frequent adjacent character\n pairs to build a subword vocabulary. Used by GPT-2, GPT-3, and RoBERTa.\n\n A word is first split into characters, then merge rules are applied in\n priority order (earlier rules have higher priority). Merging continues until\n no more rules apply.\n\n Tokenized words are cached in a direct-mapped bounded cache for amortized\n performance. *)\n\ntype t\n(** The type for BPE models. Internally mutable due to the merge cache. *)\n\ntype vocab = (string, int) Hashtbl.t\n(** The type for vocabularies mapping token strings to IDs. *)\n\ntype merges = (string * string) list\n(** The type for merge rules in priority order (earlier rules have higher\n priority). *)\n\n(** {1:creation Creation} *)\n\nval create :\n vocab:vocab ->\n merges:merges ->\n ?cache_capacity:int ->\n ?dropout:float ->\n ?unk_token:string ->\n ?continuing_subword_prefix:string ->\n ?end_of_word_suffix:string ->\n ?fuse_unk:bool ->\n ?byte_fallback:bool ->\n ?ignore_merges:bool ->\n unit ->\n t\n(** [create ~vocab ~merges ()] is a BPE model.\n\n - [cache_capacity] is the number of slots in the direct-mapped word cache.\n Defaults to [10000]. Set to [0] to disable caching. Words longer than 4096\n bytes bypass the cache.\n - [dropout] is the probability of randomly skipping a merge during\n tokenization (BPE-dropout regularization). Defaults to [0.] (no dropout).\n - [unk_token] is emitted for characters not in [vocab] (when\n {!byte_fallback} is off). No default.\n - [continuing_subword_prefix] is prepended to non-initial subwords. No\n default.\n - [end_of_word_suffix] is appended to the final subword of each word. No\n default.\n - [fuse_unk], when [true], merges consecutive unknown bytes into a single\n [unk_token] instead of emitting one per byte. Defaults to [false].\n - [byte_fallback], when [true], falls back to byte-level tokens (e.g.\n [\"<0xFF>\"]) for characters not in [vocab] instead of emitting [unk_token].\n Defaults to [false].\n - [ignore_merges], when [true], skips the merge step entirely and returns\n raw character-level tokens. Defaults to [false]. *)\n\nval from_files : vocab_file:string -> merges_file:string -> t\n(** [from_files ~vocab_file ~merges_file] loads a BPE model from\n HuggingFace-format files.\n\n - [vocab_file] is a JSON object mapping token strings to integer IDs.\n - [merges_file] is a text file with one space-separated merge pair per line.\n An optional [#version:] header line is skipped. *)\n\n(** {1:tokenization Tokenization} *)\n\ntype token = { id : int; value : string; offsets : int * int }\n(** The type for tokens. [id] is the vocabulary index, [value] the string\n content, and [offsets] the [(start, stop)] byte span in the source text. *)\n\nval tokenize : t -> string -> token list\n(** [tokenize t s] is the BPE tokenization of [s]. *)\n\nval tokenize_ids : t -> string -> int array\n(** [tokenize_ids t s] is like {!tokenize} but returns only token IDs. *)\n\nval tokenize_encoding : t -> string -> type_id:int -> Encoding.t\n(** [tokenize_encoding t s ~type_id] tokenizes [s] and builds an {!Encoding.t}\n directly, avoiding intermediate list allocation. *)\n\n(** {1:vocabulary Vocabulary} *)\n\nval token_to_id : t -> string -> int option\n(** [token_to_id t tok] is the ID of [tok] in the vocabulary. *)\n\nval id_to_token : t -> int -> string option\n(** [id_to_token t id] is the token string for [id]. *)\n\nval get_vocab : t -> (string * int) list\n(** [get_vocab t] is the vocabulary as [(token, id)] pairs. *)\n\nval get_vocab_size : t -> int\n(** [get_vocab_size t] is the number of tokens in the vocabulary. *)\n\nval get_unk_token : t -> string option\n(** [get_unk_token t] is the unknown token, if configured. *)\n\nval get_continuing_subword_prefix : t -> string option\n(** [get_continuing_subword_prefix t] is the subword prefix, if configured (e.g.\n [\"##\"]). *)\n\nval get_end_of_word_suffix : t -> string option\n(** [get_end_of_word_suffix t] is the word-end suffix, if configured (e.g.\n [\"\"]). *)\n\nval get_merges : t -> (string * string) list\n(** [get_merges t] is the merge rules in priority order. *)\n\n(** {1:serialization Serialization} *)\n\nval save : t -> path:string -> ?name:string -> unit -> unit\n(** [save t ~path ()] writes the model to [path] as two files:\n\n - [vocab.json]: a JSON object mapping token strings to IDs.\n - [merges.txt]: merge pairs, one per line, with a [#version: 0.2] header. *)\n\n(** {1:training Training} *)\n\nval train :\n min_frequency:int ->\n vocab_size:int ->\n show_progress:bool ->\n special_tokens:string list ->\n limit_alphabet:int option ->\n initial_alphabet:char list ->\n continuing_subword_prefix:string option ->\n end_of_word_suffix:string option ->\n max_token_length:int option ->\n string list ->\n t option ->\n t * string list\n(** [train ~min_frequency ~vocab_size ~show_progress ~special_tokens\n ~limit_alphabet ~initial_alphabet ~continuing_subword_prefix\n ~end_of_word_suffix ~max_token_length texts init] learns BPE merges from\n [texts].\n\n The algorithm counts word frequencies, builds an initial character alphabet,\n then iteratively finds and merges the highest-frequency adjacent pair until\n [vocab_size] is reached or pair frequency drops below [min_frequency].\n\n - [min_frequency] is the minimum pair frequency to merge.\n - [vocab_size] is the target vocabulary size.\n - [show_progress] enables progress output on [stderr].\n - [special_tokens] are added to the vocabulary first.\n - [limit_alphabet] caps the number of distinct initial characters kept.\n - [initial_alphabet] seeds the character set.\n - [continuing_subword_prefix] is set on the resulting model.\n - [end_of_word_suffix] is set on the resulting model.\n - [max_token_length] limits the byte length of merged tokens.\n - [init], when provided, seeds the vocabulary from an existing model.\n\n Returns [(model, special_tokens)]. *)\n"
},
{
"path": "packages/brot/lib/brot.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Normalizer = Normalizer\nmodule Pre_tokenizer = Pre_tokenizer\nmodule Post_processor = Post_processor\nmodule Decoder = Decoder\nmodule Encoding = Encoding\n\nlet strf = Printf.sprintf\n\n(* Error messages *)\n\nlet err_pair_no_post = \"pair sequences require a configured post-processor\"\nlet err_no_pad_token = \"padding requested but no pad token configured\"\nlet err_pad_not_in_vocab tok = strf \"pad token '%s' not in vocabulary\" tok\nlet err_add_tokens = \"only supported for word-level tokenizers\"\nlet err_export_tiktoken = \"only supported for BPE models\"\nlet err_infer_type = \"unable to infer model type from JSON\"\n\n(* Types *)\n\ntype direction = [ `Left | `Right ]\n\ntype special = {\n token : string;\n single_word : bool;\n lstrip : bool;\n rstrip : bool;\n normalized : bool;\n}\n\ntype pad_length = [ `Batch_longest | `Fixed of int | `To_multiple of int ]\n\ntype padding = {\n length : pad_length;\n direction : direction;\n pad_id : int option;\n pad_type_id : int option;\n pad_token : string option;\n}\n\ntype truncation = { max_length : int; direction : direction }\ntype data = [ `Files of string list | `Seq of string Seq.t ]\ntype sequence = { text : string; pair : string option }\n\ntype algorithm =\n | Alg_bpe of Bpe.t\n | Alg_wordpiece of Wordpiece.t\n | Alg_wordlevel of Word_level.t\n | Alg_unigram of Unigram.t\n | Alg_chars of Chars.t\n\ntype t = {\n algorithm : algorithm;\n normalizer : Normalizer.t option;\n pre_tokenizer : Pre_tokenizer.t option;\n post_processor : Post_processor.t option;\n decoder : Decoder.t option;\n specials : special list;\n special_lookup : (string, unit) Hashtbl.t;\n bos_token : string option;\n eos_token : string option;\n pad_token : string option;\n pad_id : int option;\n pad_type_id : int;\n unk_token : string option;\n}\n\nlet special ?(single_word = false) ?(lstrip = false) ?(rstrip = false)\n ?(normalized = false) token =\n { token; single_word; lstrip; rstrip; normalized }\n\nlet padding ?(direction = `Right) ?pad_id ?pad_type_id ?pad_token length =\n { length; direction; pad_id; pad_type_id; pad_token }\n\nlet truncation ?(direction = `Right) max_length = { max_length; direction }\n\n(* Algorithm dispatch *)\n\nlet alg_add_tokens algorithm tokens =\n match algorithm with\n | Alg_wordlevel model ->\n ignore (Word_level.add_tokens model tokens);\n algorithm\n | Alg_bpe _ | Alg_wordpiece _ | Alg_unigram _ | Alg_chars _ -> algorithm\n\nlet alg_token_to_id algorithm token =\n match algorithm with\n | Alg_bpe m -> Bpe.token_to_id m token\n | Alg_wordpiece m -> Wordpiece.token_to_id m token\n | Alg_wordlevel m -> Word_level.token_to_id m token\n | Alg_unigram m -> Unigram.token_to_id m token\n | Alg_chars m -> Chars.token_to_id m token\n\nlet alg_id_to_token algorithm id =\n match algorithm with\n | Alg_bpe m -> Bpe.id_to_token m id\n | Alg_wordpiece m -> Wordpiece.id_to_token m id\n | Alg_wordlevel m -> Word_level.id_to_token m id\n | Alg_unigram m -> Unigram.id_to_token m id\n | Alg_chars m -> Chars.id_to_token m id\n\nlet alg_vocab algorithm =\n match algorithm with\n | Alg_bpe m -> Bpe.get_vocab m\n | Alg_wordpiece m -> Wordpiece.get_vocab m\n | Alg_wordlevel m -> Word_level.get_vocab m\n | Alg_unigram m ->\n Unigram.get_vocab m |> List.mapi (fun i (token, _) -> (token, i))\n | Alg_chars m -> Chars.get_vocab m\n\nlet alg_vocab_size algorithm =\n match algorithm with\n | Alg_bpe m -> Bpe.get_vocab_size m\n | Alg_wordpiece m -> Wordpiece.get_vocab_size m\n | Alg_wordlevel m -> Word_level.get_vocab_size m\n | Alg_unigram m -> Unigram.get_vocab_size m\n | Alg_chars m -> Chars.get_vocab_size m\n\nlet alg_save algorithm ~folder ?prefix () =\n match algorithm with\n | Alg_bpe m ->\n Bpe.save m ~path:folder ?name:prefix ();\n let name base ext =\n match prefix with\n | Some n -> Filename.concat folder (strf \"%s-%s.%s\" n base ext)\n | None -> Filename.concat folder (strf \"%s.%s\" base ext)\n in\n [ name \"vocab\" \"json\"; name \"merges\" \"txt\" ]\n | Alg_wordpiece m -> [ Wordpiece.save m ~path:folder ?name:prefix () ]\n | Alg_wordlevel m -> Word_level.save m ~folder ()\n | Alg_unigram m -> Unigram.save m ~folder ()\n | Alg_chars m -> Chars.save m ~folder ()\n\nlet alg_tokenize algorithm text =\n match algorithm with\n | Alg_bpe m ->\n Bpe.tokenize m text\n |> List.map (fun (tok : Bpe.token) -> (tok.id, tok.value, tok.offsets))\n | Alg_wordpiece m ->\n Wordpiece.tokenize m text\n |> List.map (fun (tok : Wordpiece.token) ->\n (tok.id, tok.value, tok.offsets))\n | Alg_wordlevel m -> Word_level.tokenize m text\n | Alg_unigram m -> Unigram.tokenize m text\n | Alg_chars m -> Chars.tokenize m text\n\nlet alg_tokenize_ids algorithm text =\n match algorithm with\n | Alg_bpe m -> Bpe.tokenize_ids m text\n | Alg_wordpiece m -> Wordpiece.tokenize_ids m text\n | Alg_wordlevel m -> Word_level.tokenize_ids m text\n | Alg_unigram m ->\n Unigram.tokenize m text\n |> List.map (fun (id, _, _) -> id)\n |> Array.of_list\n | Alg_chars m ->\n Chars.tokenize m text |> List.map (fun (id, _, _) -> id) |> Array.of_list\n\nlet alg_name = function\n | Alg_bpe _ -> \"BPE\"\n | Alg_wordpiece _ -> \"WordPiece\"\n | Alg_wordlevel _ -> \"WordLevel\"\n | Alg_unigram _ -> \"Unigram\"\n | Alg_chars _ -> \"Chars\"\n\nlet vocab_to_hashtbl vocab =\n let tbl = Hashtbl.create (List.length vocab) in\n List.iter (fun (token, id) -> Hashtbl.add tbl token id) vocab;\n tbl\n\n(* Special tokens *)\n\nlet dedup_by key items =\n let seen = Hashtbl.create 16 in\n let acc = ref [] in\n List.iter\n (fun item ->\n let k = key item in\n if not (Hashtbl.mem seen k) then (\n Hashtbl.replace seen k ();\n acc := item :: !acc))\n items;\n List.rev !acc\n\nlet collect_unique_tokens specials ~bos_token ~eos_token ~pad_token ~unk_token =\n let items =\n List.map (fun (s : special) -> s.token) specials\n @ List.filter_map Fun.id [ bos_token; eos_token; pad_token; unk_token ]\n in\n dedup_by Fun.id items\n\nlet build_special_lookup specials ~bos_token ~eos_token ~pad_token ~unk_token =\n let tokens =\n collect_unique_tokens specials ~bos_token ~eos_token ~pad_token ~unk_token\n in\n let table = Hashtbl.create (List.length tokens) in\n List.iter (fun t -> Hashtbl.replace table t ()) tokens;\n table\n\n(* Construction *)\n\nlet create ?normalizer ?pre ?post ?decoder ?(specials = []) ?bos_token\n ?eos_token ?pad_token ?unk_token algorithm =\n let all_tokens =\n collect_unique_tokens specials ~bos_token ~eos_token ~pad_token ~unk_token\n in\n let algorithm = alg_add_tokens algorithm all_tokens in\n let special_lookup =\n build_special_lookup specials ~bos_token ~eos_token ~pad_token ~unk_token\n in\n let pad_id = Option.bind pad_token (alg_token_to_id algorithm) in\n {\n algorithm;\n normalizer;\n pre_tokenizer = pre;\n post_processor = post;\n decoder;\n specials;\n special_lookup;\n bos_token;\n eos_token;\n pad_token;\n pad_id;\n pad_type_id = 0;\n unk_token;\n }\n\n(* Accessors *)\n\nlet normalizer t = t.normalizer\nlet pre_tokenizer t = t.pre_tokenizer\nlet post_processor t = t.post_processor\nlet decoder t = t.decoder\nlet specials t = t.specials\nlet bos_token t = t.bos_token\nlet eos_token t = t.eos_token\nlet pad_token t = t.pad_token\nlet unk_token t = t.unk_token\nlet vocab t = alg_vocab t.algorithm\nlet vocab_size t = alg_vocab_size t.algorithm\nlet token_to_id t token = alg_token_to_id t.algorithm token\nlet id_to_token t id = alg_id_to_token t.algorithm id\n\n(* Algorithm constructors *)\n\nlet bpe ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token ?vocab ?merges ?cache_capacity ?dropout\n ?continuing_subword_prefix ?end_of_word_suffix ?fuse_unk ?byte_fallback\n ?ignore_merges () =\n let vocab_tbl =\n match vocab with None -> Hashtbl.create 100 | Some v -> vocab_to_hashtbl v\n in\n let algorithm =\n Alg_bpe\n (Bpe.create ~vocab:vocab_tbl\n ~merges:(Option.value merges ~default:[])\n ?cache_capacity ?dropout ?unk_token ?continuing_subword_prefix\n ?end_of_word_suffix ?fuse_unk ?byte_fallback ?ignore_merges ())\n in\n create ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token algorithm\n\nlet wordpiece ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token ?vocab ?continuing_subword_prefix\n ?max_input_chars_per_word () =\n let vocab_tbl =\n match vocab with None -> Hashtbl.create 100 | Some v -> vocab_to_hashtbl v\n in\n let algorithm =\n Alg_wordpiece\n (Wordpiece.create ~vocab:vocab_tbl ?unk_token ?continuing_subword_prefix\n ?max_input_chars_per_word ())\n in\n create ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token algorithm\n\nlet word_level ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token ?vocab () =\n let pre =\n match pre with Some _ -> pre | None -> Some (Pre_tokenizer.whitespace ())\n in\n let algorithm = Alg_wordlevel (Word_level.create ?vocab ?unk_token ()) in\n create ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token algorithm\n\nlet unigram ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token ?vocab () =\n let algorithm =\n Alg_unigram (Unigram.create (Option.value vocab ~default:[]))\n in\n create ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token algorithm\n\nlet chars ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token () =\n let algorithm = Alg_chars (Chars.create ()) in\n create ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token algorithm\n\nlet from_model_file ~vocab ?merges ?normalizer ?pre ?post ?decoder ?specials\n ?bos_token ?eos_token ?pad_token ?unk_token () =\n let algorithm =\n match merges with\n | Some merges_file ->\n Alg_bpe (Bpe.from_files ~vocab_file:vocab ~merges_file)\n | None -> Alg_wordpiece (Wordpiece.from_file ~vocab_file:vocab)\n in\n create ?normalizer ?pre ?post ?decoder ?specials ?bos_token ?eos_token\n ?pad_token ?unk_token algorithm\n\nlet add_tokens t tokens =\n match t.algorithm with\n | Alg_wordlevel model ->\n let vocab = Word_level.get_vocab model in\n let new_model = Word_level.create ~vocab ?unk_token:t.unk_token () in\n ignore (Word_level.add_tokens new_model tokens);\n { t with algorithm = Alg_wordlevel new_model }\n | Alg_bpe _ | Alg_wordpiece _ | Alg_unigram _ | Alg_chars _ ->\n invalid_arg err_add_tokens\n\n(* Encoding *)\n\nlet encode_text t text =\n let normalized =\n match t.normalizer with Some n -> Normalizer.apply n text | None -> text\n in\n let pre_tokens =\n match t.pre_tokenizer with\n | Some pre -> Pre_tokenizer.pre_tokenize pre normalized\n | None -> [ (normalized, (0, String.length normalized)) ]\n in\n match (t.algorithm, pre_tokens) with\n | Alg_bpe m, [ (fragment, _) ] -> Bpe.tokenize_encoding m fragment ~type_id:0\n | Alg_wordpiece m, _ ->\n Wordpiece.tokenize_spans_encoding m pre_tokens ~type_id:0\n | _ ->\n pre_tokens\n |> List.concat_map (fun (fragment, _) ->\n alg_tokenize t.algorithm fragment)\n |> Encoding.from_tokens ~type_id:0\n\nlet post_process t ~add_special primary pair =\n match t.post_processor with\n | None ->\n if Option.is_some pair then invalid_arg err_pair_no_post else primary\n | Some processor ->\n Post_processor.process processor ?pair primary\n ~add_special_tokens:add_special\n\nlet encode_single t ~add_special_tokens ~truncation seq =\n let primary = encode_text t seq.text in\n let pair = Option.map (encode_text t) seq.pair in\n let processed = post_process t ~add_special:add_special_tokens primary pair in\n match truncation with\n | None -> processed\n | Some { max_length; direction } ->\n Encoding.truncate processed ~max_length ~stride:0 ~direction\n\n(* Padding *)\n\nlet resolve_pad t (cfg : padding) =\n let token =\n match cfg.pad_token with Some _ as v -> v | None -> t.pad_token\n in\n let token =\n match token with\n | Some token -> token\n | None -> invalid_arg err_no_pad_token\n in\n let id = match cfg.pad_id with Some _ as v -> v | None -> t.pad_id in\n let id =\n match id with\n | Some id -> id\n | None -> (\n match alg_token_to_id t.algorithm token with\n | Some id -> id\n | None -> invalid_arg (err_pad_not_in_vocab token))\n in\n let type_id = Option.value cfg.pad_type_id ~default:t.pad_type_id in\n (token, id, type_id)\n\nlet round_up_to_multiple n m = if n mod m = 0 then n else (n + m - 1) / m * m\n\nlet apply_padding t encodings = function\n | None -> encodings\n | Some cfg -> (\n let pad_token, pad_id, pad_type_id = resolve_pad t cfg in\n let direction = cfg.direction in\n let pad enc target =\n if Encoding.length enc >= target then enc\n else\n Encoding.pad enc ~target_length:target ~pad_id ~pad_type_id ~pad_token\n ~direction\n in\n match cfg.length with\n | `Fixed n -> List.map (fun enc -> pad enc n) encodings\n | `Batch_longest ->\n let max_len =\n List.fold_left\n (fun acc enc -> max acc (Encoding.length enc))\n 0 encodings\n in\n List.map (fun enc -> pad enc max_len) encodings\n | `To_multiple m ->\n if m <= 0 then encodings\n else\n List.map\n (fun enc ->\n pad enc (round_up_to_multiple (Encoding.length enc) m))\n encodings)\n\n(* Parallel batch encoding *)\n\nlet encode_parallel t sequences ~add_special_tokens ~truncation =\n let arr = Array.of_list sequences in\n let n = Array.length arr in\n let results =\n Array.make n (encode_single t ~add_special_tokens ~truncation arr.(0))\n in\n let num_domains = min n (Domain.recommended_domain_count ()) in\n if num_domains <= 1 then\n for i = 1 to n - 1 do\n results.(i) <- encode_single t ~add_special_tokens ~truncation arr.(i)\n done\n else begin\n let chunk_size = n / num_domains in\n let remainder = n mod num_domains in\n let domains =\n Array.init (num_domains - 1) (fun d ->\n let start = ((d + 1) * chunk_size) + min (d + 1) remainder in\n let len = chunk_size + if d + 1 < remainder then 1 else 0 in\n Domain.spawn (fun () ->\n for i = start to start + len - 1 do\n results.(i) <-\n encode_single t ~add_special_tokens ~truncation arr.(i)\n done))\n in\n let main_len = chunk_size + if 0 < remainder then 1 else 0 in\n for i = 1 to main_len - 1 do\n results.(i) <- encode_single t ~add_special_tokens ~truncation arr.(i)\n done;\n Array.iter Domain.join domains\n end;\n Array.to_list results\n\nlet encode_sequences t sequences ~add_special_tokens ~padding ~truncation =\n let n = List.length sequences in\n let raw =\n if n >= 4 then encode_parallel t sequences ~add_special_tokens ~truncation\n else List.map (encode_single t ~add_special_tokens ~truncation) sequences\n in\n apply_padding t raw padding\n\nlet encode t ?pair ?(add_special_tokens = true) ?padding ?truncation text =\n match\n encode_sequences t\n [ { text; pair } ]\n ~add_special_tokens ~padding ~truncation\n with\n | [ encoding ] -> encoding\n | _ -> assert false\n\nlet encode_batch t ?(add_special_tokens = true) ?padding ?truncation = function\n | [] -> []\n | texts ->\n let sequences = List.map (fun text -> { text; pair = None }) texts in\n encode_sequences t sequences ~add_special_tokens ~padding ~truncation\n\nlet encode_pairs_batch t ?(add_special_tokens = true) ?padding ?truncation =\n function\n | [] -> []\n | pairs ->\n let sequences =\n List.map (fun (text, pair) -> { text; pair = Some pair }) pairs\n in\n encode_sequences t sequences ~add_special_tokens ~padding ~truncation\n\nlet encode_ids t ?pair ?add_special_tokens ?padding ?truncation text =\n let use_fast_path =\n Option.is_none pair\n && (add_special_tokens = None || add_special_tokens = Some false)\n && Option.is_none padding && Option.is_none truncation\n && Option.is_none t.post_processor\n in\n if not use_fast_path then\n Encoding.ids (encode t ?pair ?add_special_tokens ?padding ?truncation text)\n else\n let normalized =\n match t.normalizer with Some n -> Normalizer.apply n text | None -> text\n in\n let pre_tokens =\n match t.pre_tokenizer with\n | Some pre -> Pre_tokenizer.pre_tokenize pre normalized\n | None -> [ (normalized, (0, String.length normalized)) ]\n in\n let id_arrays =\n List.map\n (fun (fragment, _) -> alg_tokenize_ids t.algorithm fragment)\n pre_tokens\n in\n let total_len =\n List.fold_left (fun acc a -> acc + Array.length a) 0 id_arrays\n in\n let result = Array.make total_len 0 in\n let pos = ref 0 in\n List.iter\n (fun a ->\n let len = Array.length a in\n Array.blit a 0 result !pos len;\n pos := !pos + len)\n id_arrays;\n result\n\n(* Decoding *)\n\nlet decode t ?(skip_special_tokens = false) ids =\n let tokens =\n Array.to_list ids\n |> List.filter_map (fun id ->\n match alg_id_to_token t.algorithm id with\n | None -> None\n | Some token\n when skip_special_tokens && Hashtbl.mem t.special_lookup token ->\n None\n | Some token -> Some token)\n in\n match t.decoder with\n | Some decoder -> Decoder.decode decoder tokens\n | None -> (\n match t.algorithm with\n | Alg_wordlevel _ -> String.concat \" \" tokens\n | _ -> String.concat \"\" tokens)\n\nlet decode_batch t ?(skip_special_tokens = false) id_lists =\n List.map (decode t ~skip_special_tokens) id_lists\n\n(* Training *)\n\nlet special_tokens_for_training init specials =\n let items =\n (match specials with\n | Some sl -> List.map (fun (s : special) -> s.token) sl\n | None -> [])\n @\n match init with\n | Some tok -> List.map (fun (s : special) -> s.token) tok.specials\n | None -> []\n in\n dedup_by Fun.id items\n\nlet merge_specials_from_training ~user_specials ~trained_tokens =\n let items =\n (match user_specials with Some sl -> sl | None -> [])\n @ List.map special trained_tokens\n in\n dedup_by (fun (s : special) -> s.token) items\n\nlet data_to_strings = function\n | `Files files ->\n let lines = ref [] in\n List.iter\n (fun file ->\n let ic = open_in file in\n (try\n while true do\n lines := input_line ic :: !lines\n done\n with End_of_file -> ());\n close_in ic)\n files;\n List.rev !lines\n | `Seq seq -> List.of_seq seq\n\nlet initial_alphabet_of strs =\n List.map (fun s -> if String.length s > 0 then s.[0] else ' ') strs\n\nlet train_bpe ?init ?normalizer ?pre ?post ?decoder ?specials ?bos_token\n ?eos_token ?pad_token ?unk_token ?(vocab_size = 30000) ?(min_frequency = 0)\n ?limit_alphabet ?initial_alphabet ?continuing_subword_prefix\n ?end_of_word_suffix ?(show_progress = true) ?max_token_length data =\n let special_tokens = special_tokens_for_training init specials in\n let initial_alphabet =\n Option.value initial_alphabet ~default:[] |> initial_alphabet_of\n in\n let limit_alphabet = Some (Option.value limit_alphabet ~default:1000) in\n let texts = data_to_strings data in\n let existing_bpe =\n Option.bind init (fun t ->\n match t.algorithm with Alg_bpe m -> Some m | _ -> None)\n in\n let trained_model, result_specials =\n Bpe.train ~min_frequency ~vocab_size ~show_progress ~special_tokens\n ~limit_alphabet ~initial_alphabet ~continuing_subword_prefix\n ~end_of_word_suffix ~max_token_length texts existing_bpe\n in\n let all_specials =\n merge_specials_from_training ~user_specials:specials\n ~trained_tokens:result_specials\n in\n create ?normalizer ?pre ?post ?decoder ~specials:all_specials ?bos_token\n ?eos_token ?pad_token ?unk_token (Alg_bpe trained_model)\n\nlet train_wordpiece ?init ?normalizer ?pre ?post ?decoder ?specials ?bos_token\n ?eos_token ?pad_token ?unk_token ?(vocab_size = 30000) ?(min_frequency = 0)\n ?limit_alphabet ?initial_alphabet ?(continuing_subword_prefix = \"##\")\n ?end_of_word_suffix ?(show_progress = true) data =\n let special_tokens = special_tokens_for_training init specials in\n let initial_alphabet =\n Option.value initial_alphabet ~default:[] |> initial_alphabet_of\n in\n let limit_alphabet = Some (Option.value limit_alphabet ~default:1000) in\n let texts = data_to_strings data in\n let existing_wp =\n Option.bind init (fun t ->\n match t.algorithm with Alg_wordpiece m -> Some m | _ -> None)\n in\n let trained_model, result_specials =\n Wordpiece.train ~min_frequency ~vocab_size ~show_progress ~special_tokens\n ~limit_alphabet ~initial_alphabet ~continuing_subword_prefix\n ~end_of_word_suffix texts existing_wp\n in\n let all_specials =\n merge_specials_from_training ~user_specials:specials\n ~trained_tokens:result_specials\n in\n create ?normalizer ?pre ?post ?decoder ~specials:all_specials ?bos_token\n ?eos_token ?pad_token ?unk_token (Alg_wordpiece trained_model)\n\nlet train_wordlevel ?init ?normalizer ?pre ?post ?decoder ?specials ?bos_token\n ?eos_token ?pad_token ?unk_token ?(vocab_size = 30000) ?(min_frequency = 0)\n ?(show_progress = true) data =\n let special_tokens = special_tokens_for_training init specials in\n let texts = data_to_strings data in\n let existing_wl =\n Option.bind init (fun t ->\n match t.algorithm with Alg_wordlevel m -> Some m | _ -> None)\n in\n let trained_model, result_specials =\n Word_level.train ~vocab_size ~min_frequency ~show_progress ~special_tokens\n texts existing_wl\n in\n let all_specials =\n merge_specials_from_training ~user_specials:specials\n ~trained_tokens:result_specials\n in\n create ?normalizer ?pre ?post ?decoder ~specials:all_specials ?bos_token\n ?eos_token ?pad_token ?unk_token (Alg_wordlevel trained_model)\n\nlet train_unigram ?init ?normalizer ?pre ?post ?decoder ?specials ?bos_token\n ?eos_token ?pad_token ?unk_token ?(vocab_size = 8000)\n ?(show_progress = true) ?(shrinking_factor = 0.75) ?(max_piece_length = 16)\n ?(n_sub_iterations = 2) data =\n let special_tokens = special_tokens_for_training init specials in\n let texts = data_to_strings data in\n let existing_ug =\n Option.bind init (fun t ->\n match t.algorithm with Alg_unigram m -> Some m | _ -> None)\n in\n let trained_model, result_specials =\n Unigram.train ~vocab_size ~show_progress ~special_tokens ~shrinking_factor\n ~unk_token ~max_piece_length ~n_sub_iterations texts existing_ug\n in\n let all_specials =\n merge_specials_from_training ~user_specials:specials\n ~trained_tokens:result_specials\n in\n create ?normalizer ?pre ?post ?decoder ~specials:all_specials ?bos_token\n ?eos_token ?pad_token ?unk_token (Alg_unigram trained_model)\n\n(* JSON serialization *)\n\nlet json_obj pairs =\n Jsont.Json.object' (List.map (fun (k, v) -> (Jsont.Json.name k, v)) pairs)\n\nlet json_mem name = function\n | Jsont.Object (mems, _) -> (\n match Jsont.Json.find_mem name mems with\n | Some (_, v) -> v\n | None -> Jsont.Null ((), Jsont.Meta.none))\n | _ -> Jsont.Null ((), Jsont.Meta.none)\n\nlet json_string_or_null = function Jsont.String (s, _) -> Some s | _ -> None\nlet json_option_of f = function None -> Jsont.Json.null () | Some v -> f v\n\nlet special_of_json json =\n let mem name = json_mem name json in\n let to_bool = function Jsont.Bool (b, _) -> b | _ -> false in\n let to_str = function\n | Jsont.String (s, _) -> s\n | _ -> failwith \"expected string\"\n in\n {\n token = to_str (mem \"content\");\n single_word = to_bool (mem \"single_word\");\n lstrip = to_bool (mem \"lstrip\");\n rstrip = to_bool (mem \"rstrip\");\n normalized = to_bool (mem \"normalized\");\n }\n\nlet added_token_to_json ~id (s : special) =\n json_obj\n [\n (\"id\", Jsont.Json.int id);\n (\"content\", Jsont.Json.string s.token);\n (\"single_word\", Jsont.Json.bool s.single_word);\n (\"lstrip\", Jsont.Json.bool s.lstrip);\n (\"rstrip\", Jsont.Json.bool s.rstrip);\n (\"normalized\", Jsont.Json.bool s.normalized);\n (\"special\", Jsont.Json.bool true);\n ]\n\nlet vocab_to_json vocab =\n json_obj (List.map (fun (token, id) -> (token, Jsont.Json.int id)) vocab)\n\nlet alg_to_json = function\n | Alg_bpe bpe ->\n let vocab_json = vocab_to_json (Bpe.get_vocab bpe) in\n let merges_json =\n Bpe.get_merges bpe\n |> List.map (fun (a, b) ->\n Jsont.Json.list [ Jsont.Json.string a; Jsont.Json.string b ])\n |> Jsont.Json.list\n in\n json_obj\n [\n (\"type\", Jsont.Json.string \"BPE\");\n (\"dropout\", Jsont.Json.null ());\n (\"unk_token\", json_option_of Jsont.Json.string (Bpe.get_unk_token bpe));\n ( \"continuing_subword_prefix\",\n json_option_of Jsont.Json.string\n (Bpe.get_continuing_subword_prefix bpe) );\n ( \"end_of_word_suffix\",\n json_option_of Jsont.Json.string (Bpe.get_end_of_word_suffix bpe) );\n (\"fuse_unk\", Jsont.Json.bool false);\n (\"byte_fallback\", Jsont.Json.bool false);\n (\"ignore_merges\", Jsont.Json.bool false);\n (\"vocab\", vocab_json);\n (\"merges\", merges_json);\n ]\n | Alg_wordpiece wp ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"WordPiece\");\n (\"unk_token\", Jsont.Json.string (Wordpiece.get_unk_token wp));\n ( \"continuing_subword_prefix\",\n Jsont.Json.string (Wordpiece.get_continuing_subword_prefix wp) );\n (\"max_input_chars_per_word\", Jsont.Json.int 100);\n (\"vocab\", vocab_to_json (Wordpiece.get_vocab wp));\n ]\n | Alg_wordlevel wl ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"WordLevel\");\n (\"unk_token\", Jsont.Json.string \"[UNK]\");\n (\"vocab\", vocab_to_json (Word_level.get_vocab wl));\n ]\n | Alg_unigram ug ->\n let vocab_json =\n Unigram.get_vocab ug\n |> List.map (fun (token, score) ->\n Jsont.Json.list [ Jsont.Json.string token; Jsont.Json.number score ])\n |> Jsont.Json.list\n in\n json_obj\n [\n (\"type\", Jsont.Json.string \"Unigram\");\n (\"unk_id\", Jsont.Json.null ());\n (\"vocab\", vocab_json);\n ]\n | Alg_chars _ ->\n json_obj [ (\"type\", Jsont.Json.string \"Chars\"); (\"vocab\", json_obj []) ]\n\nlet to_json (t : t) =\n let vocab_list = alg_vocab t.algorithm in\n let added_tokens =\n t.specials\n |> List.filter_map (fun spec ->\n List.find_opt (fun (token, _) -> token = spec.token) vocab_list\n |> Option.map (fun (_, id) -> added_token_to_json ~id spec))\n in\n json_obj\n [\n (\"version\", Jsont.Json.string \"1.0\");\n (\"truncation\", Jsont.Json.null ());\n (\"padding\", Jsont.Json.null ());\n (\"added_tokens\", Jsont.Json.list added_tokens);\n (\"normalizer\", json_option_of Normalizer.to_json t.normalizer);\n (\"pre_tokenizer\", json_option_of Pre_tokenizer.to_json t.pre_tokenizer);\n (\"post_processor\", json_option_of Post_processor.to_json t.post_processor);\n (\"decoder\", json_option_of Decoder.to_json t.decoder);\n (\"model\", alg_to_json t.algorithm);\n ]\n\n(* JSON deserialization helpers *)\n\nlet json_to_assoc = function\n | Jsont.Object (mems, _) ->\n List.map\n (fun ((k, _), v) ->\n match v with\n | Jsont.Number (f, _) -> (k, int_of_float f)\n | _ -> failwith (\"Expected number for vocab entry: \" ^ k))\n mems\n | _ -> failwith \"Expected object for vocab\"\n\nlet json_to_list = function\n | Jsont.Array (l, _) -> l\n | _ -> failwith \"Expected array\"\n\nlet json_to_string = function\n | Jsont.String (s, _) -> s\n | _ -> failwith \"Expected string\"\n\nlet json_to_float = function\n | Jsont.Number (f, _) -> f\n | _ -> failwith \"Expected number\"\n\nlet json_has_field name j =\n match json_mem name j with Jsont.Null _ -> false | _ -> true\n\nlet json_result_to_option of_json = function\n | Jsont.Null _ -> None\n | j -> ( match of_json j with Ok v -> Some v | Error msg -> failwith msg)\n\nlet infer_model_type mj =\n match json_string_or_null (json_mem \"type\" mj) with\n | Some s -> s\n | None ->\n if json_has_field \"merges\" mj then \"BPE\"\n else if json_has_field \"unk_id\" mj then \"Unigram\"\n else if\n json_has_field \"continuing_subword_prefix\" mj\n || json_has_field \"max_input_chars_per_word\" mj\n then \"WordPiece\"\n else if json_has_field \"vocab\" mj then \"WordLevel\"\n else failwith err_infer_type\n\nlet parse_merge = function\n | Jsont.Array ([ a; b ], _) -> (json_to_string a, json_to_string b)\n | Jsont.String (s, _) -> (\n match String.split_on_char ' ' s with\n | [ a; b ] -> (a, b)\n | _ -> failwith \"Invalid merge string format\")\n | _ -> failwith \"Invalid merge entry\"\n\nlet alg_of_json mj =\n let mem name = json_mem name mj in\n let str name = json_string_or_null (mem name) in\n match infer_model_type mj with\n | \"BPE\" ->\n let vocab_list = json_to_assoc (mem \"vocab\") in\n let merges = json_to_list (mem \"merges\") |> List.map parse_merge in\n Alg_bpe\n (Bpe.create\n ~vocab:(vocab_to_hashtbl vocab_list)\n ~merges ?unk_token:(str \"unk_token\")\n ?continuing_subword_prefix:(str \"continuing_subword_prefix\")\n ?end_of_word_suffix:(str \"end_of_word_suffix\") ())\n | \"WordPiece\" ->\n let vocab_list = json_to_assoc (mem \"vocab\") in\n let unk_token = str \"unk_token\" |> Option.value ~default:\"[UNK]\" in\n let continuing_subword_prefix =\n str \"continuing_subword_prefix\" |> Option.value ~default:\"##\"\n in\n let max_input_chars_per_word =\n match mem \"max_input_chars_per_word\" with\n | Jsont.Number (f, _) -> int_of_float f\n | _ -> 100\n in\n Alg_wordpiece\n (Wordpiece.create\n ~vocab:(vocab_to_hashtbl vocab_list)\n ~unk_token ~continuing_subword_prefix ~max_input_chars_per_word ())\n | \"WordLevel\" ->\n let vocab_list = json_to_assoc (mem \"vocab\") in\n let unk_token = str \"unk_token\" |> Option.value ~default:\"[UNK]\" in\n Alg_wordlevel (Word_level.create ~vocab:vocab_list ~unk_token ())\n | \"Unigram\" ->\n let vocab =\n json_to_list (mem \"vocab\")\n |> List.map (fun arr ->\n match json_to_list arr with\n | [ token; score ] -> (json_to_string token, json_to_float score)\n | _ -> failwith \"Invalid unigram vocab format\")\n in\n Alg_unigram (Unigram.create vocab)\n | \"Chars\" -> Alg_chars (Chars.create ())\n | s -> failwith (strf \"Unsupported model type: %s\" s)\n\nlet from_json json =\n try\n let mem name = json_mem name json in\n let normalizer =\n json_result_to_option Normalizer.of_json (mem \"normalizer\")\n in\n let pre =\n json_result_to_option Pre_tokenizer.of_json (mem \"pre_tokenizer\")\n in\n let post =\n json_result_to_option Post_processor.of_json (mem \"post_processor\")\n in\n let decoder = json_result_to_option Decoder.of_json (mem \"decoder\") in\n let algorithm = alg_of_json (mem \"model\") in\n let added_tokens =\n match mem \"added_tokens\" with\n | Jsont.Array (l, _) -> List.map special_of_json l\n | _ -> []\n in\n Ok (create ?normalizer ?pre ?post ?decoder ~specials:added_tokens algorithm)\n with\n | Failure msg -> Error msg\n | exn -> Error (Printexc.to_string exn)\n\n(* File I/O *)\n\nlet write_string_to_file path s =\n let oc = open_out path in\n Fun.protect ~finally:(fun () -> close_out oc) (fun () -> output_string oc s)\n\nlet from_file path =\n try\n let ic = open_in path in\n let s =\n Fun.protect\n ~finally:(fun () -> close_in ic)\n (fun () -> really_input_string ic (in_channel_length ic))\n in\n match Jsont_bytesrw.decode_string Jsont.json s with\n | Ok json -> from_json json\n | Error e -> Error e\n with\n | Sys_error msg -> Error (\"File error: \" ^ msg)\n | exn -> Error (Printexc.to_string exn)\n\nlet save_pretrained t ~path =\n (try Sys.mkdir path 0o755 with Sys_error _ -> ());\n let json_str =\n match\n Jsont_bytesrw.encode_string ~format:Jsont.Minify Jsont.json (to_json t)\n with\n | Ok s -> s\n | Error e -> failwith (\"save_pretrained: failed to encode JSON: \" ^ e)\n in\n write_string_to_file (Filename.concat path \"tokenizer.json\") json_str\n\nlet export_tiktoken t ~merges_path ~vocab_path =\n match t.algorithm with\n | Alg_bpe bpe ->\n let vocab =\n alg_vocab t.algorithm\n |> List.sort (fun (_, id1) (_, id2) -> Int.compare id1 id2)\n in\n let json_str =\n match\n Jsont_bytesrw.encode_string ~format:Jsont.Minify Jsont.json\n (vocab_to_json vocab)\n with\n | Ok s -> s\n | Error e -> failwith (\"export_tiktoken: failed to encode vocab: \" ^ e)\n in\n write_string_to_file vocab_path json_str;\n let oc = open_out merges_path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () ->\n output_string oc \"#version: 0.2\\n\";\n List.iter\n (fun (a, b) -> Printf.fprintf oc \"%s %s\\n\" a b)\n (Bpe.get_merges bpe))\n | _ -> invalid_arg err_export_tiktoken\n\nlet save_model_files t ~folder ?prefix () =\n alg_save t.algorithm ~folder ?prefix ()\n\n(* Formatting *)\n\nlet pp ppf t =\n let yes_no = function Some _ -> \"yes\" | None -> \"no\" in\n Format.fprintf ppf\n \"@[<1>@]\"\n (alg_name t.algorithm)\n (alg_vocab_size t.algorithm)\n (yes_no t.normalizer) (yes_no t.pre_tokenizer) (yes_no t.post_processor)\n (yes_no t.decoder)\n"
},
{
"path": "packages/brot/lib/brot.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Tokenization for OCaml.\n\n Brot tokenizes text into token IDs for language models and reverses the\n process. Tokenization proceeds through configurable stages:\n\n + {e Normalization}: clean and normalize text (lowercase, accent removal,\n Unicode normalization). See {!Normalizer}.\n + {e Pre-tokenization}: split text into words or sub-words. See\n {!Pre_tokenizer}.\n + {e Tokenization}: apply vocabulary-based encoding (BPE, WordPiece,\n Unigram, word-level, or character-level).\n + {e Post-processing}: add special tokens and set type IDs. See\n {!Post_processor}.\n + {e Padding/Truncation}: adjust sequence lengths for batching.\n\n Each stage is optional and configurable. Open the module to use it, it\n defines only modules in your scope.\n\n {1:quick_start Quick start}\n\n Load a pretrained tokenizer:\n {[\n let tokenizer = Brot.from_file \"tokenizer.json\" |> Result.get_ok in\n let encoding = Brot.encode tokenizer \"Hello world!\" in\n let _ids = Encoding.ids encoding\n ]}\n\n Create a BPE tokenizer from scratch:\n {[\n let tokenizer =\n Brot.bpe\n ~vocab:[(\"hello\", 0); (\"world\", 1); (\"[PAD]\", 2)]\n ~merges:[]\n ()\n in\n let encoding = Brot.encode tokenizer \"hello world\" in\n let _text = Brot.decode tokenizer (Encoding.ids encoding)\n ]}\n\n Train a new tokenizer:\n {[\n let texts = [ \"Hello world\"; \"How are you?\"; \"Hello again\" ] in\n let tokenizer =\n Brot.train_bpe (`Seq (List.to_seq texts)) ~vocab_size:1000\n in\n Brot.save_pretrained tokenizer ~path:\"./my_tokenizer\"\n ]}\n\n {!modules:Encoding Normalizer Pre_tokenizer Post_processor Decoder} *)\n\nmodule Normalizer = Normalizer\n(** Text normalization. *)\n\nmodule Pre_tokenizer = Pre_tokenizer\n(** Pre-tokenization. *)\n\nmodule Post_processor = Post_processor\n(** Post-processing. *)\n\nmodule Decoder = Decoder\n(** Token decoding. *)\n\nmodule Encoding = Encoding\n(** Tokenization encodings. *)\n\n(** {1:types Types} *)\n\ntype t\n(** The type for tokenizers. Immutable after creation. *)\n\ntype direction = [ `Left | `Right ]\n(** The type for padding and truncation directions. [`Left] operates at the\n beginning of the sequence, [`Right] at the end. *)\n\ntype special = {\n token : string; (** The token text (e.g., [\"\"], [\"\"]). *)\n single_word : bool; (** Whether this token must match whole words only. *)\n lstrip : bool; (** Whether to strip whitespace on the left. *)\n rstrip : bool; (** Whether to strip whitespace on the right. *)\n normalized : bool; (** Whether to apply normalization to this token. *)\n}\n(** The type for special token configurations.\n\n Special tokens are never split during tokenization and can be skipped during\n decoding. Token IDs are assigned automatically when added to the vocabulary.\n The semantic role (pad, unk, bos, etc.) is contextual, not encoded in the\n type. *)\n\ntype pad_length = [ `Batch_longest | `Fixed of int | `To_multiple of int ]\n(** The type for padding length strategies.\n\n - [`Batch_longest]: pad to the longest sequence in the batch.\n - [`Fixed n]: pad every sequence to exactly [n] tokens.\n - [`To_multiple n]: pad to the smallest multiple of [n] that is at least the\n sequence length. *)\n\ntype padding = {\n length : pad_length;\n direction : direction;\n pad_id : int option;\n pad_type_id : int option;\n pad_token : string option;\n}\n(** The type for padding configurations.\n\n When [pad_id], [pad_type_id], or [pad_token] are [None], the tokenizer's\n configured padding token is used. Raises [Invalid_argument] at padding time\n if no padding token is configured and these fields are [None]. *)\n\ntype truncation = { max_length : int; direction : direction }\n(** The type for truncation configurations. Sequences exceeding [max_length]\n tokens are trimmed from the given [direction]. *)\n\ntype data = [ `Files of string list | `Seq of string Seq.t ]\n(** The type for training data sources.\n\n - [`Files paths]: read training text from files, one line per example.\n - [`Seq seq]: use a sequence of strings. *)\n\nval special :\n ?single_word:bool ->\n ?lstrip:bool ->\n ?rstrip:bool ->\n ?normalized:bool ->\n string ->\n special\n(** [special token] is a special token configuration for [token].\n\n [single_word] defaults to [false]. [lstrip] and [rstrip] default to [false].\n [normalized] defaults to [false]. *)\n\nval padding :\n ?direction:direction ->\n ?pad_id:int ->\n ?pad_type_id:int ->\n ?pad_token:string ->\n pad_length ->\n padding\n(** [padding length] is a padding configuration for the given [length] strategy.\n\n [direction] defaults to [`Right]. Other fields default to [None] (falls back\n to the tokenizer's configured padding token). *)\n\nval truncation : ?direction:direction -> int -> truncation\n(** [truncation max_length] is a truncation configuration limiting sequences to\n [max_length] tokens. [direction] defaults to [`Right]. *)\n\n(** {1:constructors Constructors} *)\n\nval bpe :\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n ?vocab:(string * int) list ->\n ?merges:(string * string) list ->\n ?cache_capacity:int ->\n ?dropout:float ->\n ?continuing_subword_prefix:string ->\n ?end_of_word_suffix:string ->\n ?fuse_unk:bool ->\n ?byte_fallback:bool ->\n ?ignore_merges:bool ->\n unit ->\n t\n(** [bpe ()] is a BPE (Byte Pair Encoding) tokenizer. Used by GPT-2, GPT-3,\n RoBERTa.\n\n - [normalizer]: text normalization. Default: none.\n - [pre]: pre-tokenization strategy. Default: none.\n - [post]: post-processor for special tokens. Default: none.\n - [decoder]: decoding strategy. Default: none.\n - [specials]: special tokens to add to vocabulary. Default: [[]].\n - [bos_token], [eos_token], [pad_token]: role markers; added to vocabulary\n if not already present. Default: none.\n - [unk_token]: token for unknown characters. Configures both the role and\n the BPE model's unknown handling. Default: none.\n - [vocab]: initial vocabulary as [(token, id)] pairs. Default: [[]].\n - [merges]: merge rules as [(left, right)] pairs learned during training.\n Default: [[]].\n - [cache_capacity]: LRU cache size for tokenization results. Default:\n [10000].\n - [dropout]: probability \\[[0]; [1]\\] of skipping merges (data\n augmentation). Default: none (no dropout).\n - [continuing_subword_prefix]: prefix for non-initial subwords (e.g.,\n [\"##\"]). Default: none.\n - [end_of_word_suffix]: suffix marking word boundaries (e.g., [\"\"]).\n Default: none.\n - [fuse_unk]: merge consecutive unknown tokens. Default: [false].\n - [byte_fallback]: use byte-level fallback ([\"<0x00>\"]) instead of unknown\n token. Default: [false].\n - [ignore_merges]: skip merge application (character-level output). Default:\n [false]. *)\n\nval wordpiece :\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n ?vocab:(string * int) list ->\n ?continuing_subword_prefix:string ->\n ?max_input_chars_per_word:int ->\n unit ->\n t\n(** [wordpiece ()] is a WordPiece tokenizer. Used by BERT, DistilBERT, Electra.\n\n WordPiece uses a greedy longest-match-first algorithm to split words into\n subword pieces prefixed with a continuation marker (e.g., [\"running\"]\n becomes [[\"run\"; \"##ning\"]]).\n\n - [vocab]: initial vocabulary as [(token, id)] pairs. Default: [[]].\n - [unk_token]: token for out-of-vocabulary words. Default: [\"[UNK]\"].\n - [continuing_subword_prefix]: prefix for non-initial subwords. Default:\n [\"##\"].\n - [max_input_chars_per_word]: words longer than this are replaced with\n [unk_token]. Default: [100].\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token]) are as in {!bpe}. *)\n\nval word_level :\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n ?vocab:(string * int) list ->\n unit ->\n t\n(** [word_level ()] is a word-level tokenizer.\n\n Maps each word directly to a token ID. No subword splitting is performed.\n Words not in vocabulary map to [unk_token].\n\n {b Note.} When [pre] is not provided, {!Pre_tokenizer.whitespace} is used by\n default.\n\n - [vocab]: initial vocabulary as [(word, id)] pairs. Default: [[]].\n - [unk_token]: token for out-of-vocabulary words. Default: [\"\"].\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token]) are as in {!bpe}. *)\n\nval unigram :\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n ?vocab:(string * float) list ->\n unit ->\n t\n(** [unigram ()] is a Unigram tokenizer. Used by AlBERT, T5, mBART.\n\n Unigram uses probabilistic segmentation to find optimal subword splits based\n on token log-probabilities.\n\n - [vocab]: initial vocabulary as [(token, score)] pairs where scores are\n negative log probabilities. Default: [[]].\n - [unk_token]: token for unknown characters. Default: none.\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token]) are as in {!bpe}. *)\n\nval chars :\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n unit ->\n t\n(** [chars ()] is a character-level tokenizer.\n\n Each byte in the input becomes a separate token with ID equal to its ordinal\n value. No vocabulary is required.\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token]) are as in {!bpe}. *)\n\nval from_model_file :\n vocab:string ->\n ?merges:string ->\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n unit ->\n t\n(** [from_model_file ~vocab ()] loads a tokenizer from HuggingFace model files.\n\n The model type is inferred from the arguments: if [merges] is provided, a\n BPE tokenizer is created; otherwise WordPiece.\n\n - [vocab]: path to vocabulary file ([vocab.json]). Expected format: JSON\n object mapping tokens to IDs ([{\"hello\": 0, \"world\": 1}]).\n - [merges]: path to merges file ([merges.txt]). One merge per line as\n space-separated token pairs. Lines starting with [\"#version\"] are skipped.\n\n Raises [Sys_error] if a file cannot be read.\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token], [unk_token]) are as in {!bpe}. *)\n\nval add_tokens : t -> string list -> t\n(** [add_tokens t tokens] is [t] with [tokens] added to the vocabulary. Only\n supported for word-level tokenizers.\n\n Raises [Invalid_argument] if the tokenizer does not support dynamic\n vocabulary extension. *)\n\n(** {1:accessors Accessors} *)\n\nval normalizer : t -> Normalizer.t option\n(** [normalizer t] is [t]'s normalizer, if any. *)\n\nval pre_tokenizer : t -> Pre_tokenizer.t option\n(** [pre_tokenizer t] is [t]'s pre-tokenizer, if any. *)\n\nval post_processor : t -> Post_processor.t option\n(** [post_processor t] is [t]'s post-processor, if any. *)\n\nval decoder : t -> Decoder.t option\n(** [decoder t] is [t]'s decoder, if any. *)\n\nval specials : t -> special list\n(** [specials t] is [t]'s special tokens. *)\n\nval bos_token : t -> string option\n(** [bos_token t] is [t]'s beginning-of-sequence token, if any. *)\n\nval eos_token : t -> string option\n(** [eos_token t] is [t]'s end-of-sequence token, if any. *)\n\nval pad_token : t -> string option\n(** [pad_token t] is [t]'s padding token, if any. *)\n\nval unk_token : t -> string option\n(** [unk_token t] is [t]'s unknown token, if any. *)\n\n(** {1:vocab Vocabulary} *)\n\nval vocab : t -> (string * int) list\n(** [vocab t] is [t]'s vocabulary as [(token, id)] pairs. *)\n\nval vocab_size : t -> int\n(** [vocab_size t] is the number of tokens in [t]'s vocabulary. *)\n\nval token_to_id : t -> string -> int option\n(** [token_to_id t token] is the ID of [token] in [t], if any. *)\n\nval id_to_token : t -> int -> string option\n(** [id_to_token t id] is the token string for [id] in [t], if any. *)\n\n(** {1:encoding Encoding and decoding} *)\n\nval encode :\n t ->\n ?pair:string ->\n ?add_special_tokens:bool ->\n ?padding:padding ->\n ?truncation:truncation ->\n string ->\n Encoding.t\n(** [encode t text] is the encoding of [text] by [t].\n\n - [pair]: a second sentence for sentence-pair tasks. The post-processor\n merges both sequences with appropriate type IDs. Default: none.\n - [add_special_tokens]: whether to insert special tokens via the\n post-processor. Default: [true].\n - [padding]: padding configuration. Default: none (no padding).\n - [truncation]: truncation configuration. Default: none (no truncation). *)\n\nval encode_batch :\n t ->\n ?add_special_tokens:bool ->\n ?padding:padding ->\n ?truncation:truncation ->\n string list ->\n Encoding.t list\n(** [encode_batch t texts] is the encoding of each text in [texts].\n\n Optional parameters are as in {!encode}. For sentence-pair tasks, use\n {!encode_pairs_batch}. *)\n\nval encode_pairs_batch :\n t ->\n ?add_special_tokens:bool ->\n ?padding:padding ->\n ?truncation:truncation ->\n (string * string) list ->\n Encoding.t list\n(** [encode_pairs_batch t pairs] encodes a batch of sentence pairs. Each element\n is [(primary, secondary)].\n\n Optional parameters are as in {!encode}. *)\n\nval encode_ids :\n t ->\n ?pair:string ->\n ?add_special_tokens:bool ->\n ?padding:padding ->\n ?truncation:truncation ->\n string ->\n int array\n(** [encode_ids t text] is [Encoding.ids (encode t text)].\n\n Optional parameters are as in {!encode}. *)\n\nval decode : t -> ?skip_special_tokens:bool -> int array -> string\n(** [decode t ids] is the text obtained by decoding [ids] through [t]'s\n vocabulary and decoder.\n\n [skip_special_tokens] defaults to [false]. *)\n\nval decode_batch :\n t -> ?skip_special_tokens:bool -> int array list -> string list\n(** [decode_batch t ids_list] decodes each element of [ids_list].\n\n [skip_special_tokens] defaults to [false]. *)\n\n(** {1:training Training} *)\n\nval train_bpe :\n ?init:t ->\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n ?vocab_size:int ->\n ?min_frequency:int ->\n ?limit_alphabet:int ->\n ?initial_alphabet:string list ->\n ?continuing_subword_prefix:string ->\n ?end_of_word_suffix:string ->\n ?show_progress:bool ->\n ?max_token_length:int ->\n data ->\n t\n(** [train_bpe data] trains a BPE tokenizer from [data].\n\n Learns merge rules by iteratively merging the most frequent adjacent pairs\n until reaching the target vocabulary size.\n\n - [init]: existing tokenizer to extend. Default: create new.\n - [vocab_size]: target vocabulary size including special tokens. Default:\n [30000].\n - [min_frequency]: minimum pair frequency to be merged. Default: [0].\n - [limit_alphabet]: maximum number of initial characters to keep. Default:\n none (keep all).\n - [initial_alphabet]: characters to include regardless of frequency.\n Default: [[]].\n - [continuing_subword_prefix]: prefix for non-initial subwords. Default:\n none.\n - [end_of_word_suffix]: suffix marking word boundaries. Default: none.\n - [show_progress]: display progress bar. Default: [true].\n - [max_token_length]: maximum token length. Default: none.\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token], [unk_token]) are as in {!bpe}. *)\n\nval train_wordpiece :\n ?init:t ->\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n ?vocab_size:int ->\n ?min_frequency:int ->\n ?limit_alphabet:int ->\n ?initial_alphabet:string list ->\n ?continuing_subword_prefix:string ->\n ?end_of_word_suffix:string ->\n ?show_progress:bool ->\n data ->\n t\n(** [train_wordpiece data] trains a WordPiece tokenizer from [data].\n\n Learns subword vocabulary by maximizing language model likelihood.\n\n - [init]: existing tokenizer to extend. Default: create new.\n - [vocab_size]: target vocabulary size including special tokens. Default:\n [30000].\n - [min_frequency]: minimum frequency for a subword to be included. Default:\n [0].\n - [limit_alphabet]: maximum number of initial characters to keep. Default:\n none (keep all).\n - [initial_alphabet]: characters to include regardless of frequency.\n Default: [[]].\n - [continuing_subword_prefix]: prefix for non-initial subwords. Default:\n [\"##\"].\n - [end_of_word_suffix]: suffix marking word boundaries. Default: none.\n - [show_progress]: display progress bar. Default: [true].\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token], [unk_token]) are as in {!bpe}. *)\n\nval train_wordlevel :\n ?init:t ->\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n ?vocab_size:int ->\n ?min_frequency:int ->\n ?show_progress:bool ->\n data ->\n t\n(** [train_wordlevel data] trains a word-level tokenizer from [data].\n\n Builds vocabulary by collecting unique words, optionally filtering by\n frequency. No subword splitting.\n\n - [init]: existing tokenizer to extend. Default: create new.\n - [vocab_size]: target vocabulary size including special tokens. Default:\n [30000].\n - [min_frequency]: minimum frequency for a word to be included. Default:\n [0].\n - [show_progress]: display progress bar. Default: [true].\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token], [unk_token]) are as in {!bpe}. *)\n\nval train_unigram :\n ?init:t ->\n ?normalizer:Normalizer.t ->\n ?pre:Pre_tokenizer.t ->\n ?post:Post_processor.t ->\n ?decoder:Decoder.t ->\n ?specials:special list ->\n ?bos_token:string ->\n ?eos_token:string ->\n ?pad_token:string ->\n ?unk_token:string ->\n ?vocab_size:int ->\n ?show_progress:bool ->\n ?shrinking_factor:float ->\n ?max_piece_length:int ->\n ?n_sub_iterations:int ->\n data ->\n t\n(** [train_unigram data] trains a Unigram tokenizer from [data].\n\n Learns probabilistic subword vocabulary using EM algorithm.\n\n - [init]: existing tokenizer to extend. Default: create new.\n - [vocab_size]: target vocabulary size including special tokens. Default:\n [8000].\n - [show_progress]: display progress bar. Default: [true].\n - [shrinking_factor]: fraction of vocabulary to retain in each pruning\n iteration. Default: [0.75].\n - [max_piece_length]: maximum subword length. Default: [16].\n - [n_sub_iterations]: number of EM sub-iterations per pruning round.\n Default: [2].\n\n Pipeline parameters ([normalizer], [pre], [post], [decoder], [specials],\n [bos_token], [eos_token], [pad_token], [unk_token]) are as in {!bpe}. *)\n\n(** {1:model_files Model files} *)\n\nval export_tiktoken : t -> merges_path:string -> vocab_path:string -> unit\n(** [export_tiktoken t ~merges_path ~vocab_path] exports [t]'s BPE merges and\n vocabulary in tiktoken-compatible format.\n\n {b Warning.} Only BPE tokenizers are supported. Raises [Failure] for other\n model types. *)\n\nval save_model_files :\n t -> folder:string -> ?prefix:string -> unit -> string list\n(** [save_model_files t ~folder ?prefix ()] saves [t]'s underlying model files\n (vocabulary and merges) to [folder] and returns the list of created file\n paths.\n\n [prefix] defaults to [\"\"]. *)\n\n(** {1:huggingface HuggingFace compatibility} *)\n\nval from_file : string -> (t, string) result\n(** [from_file path] is a tokenizer loaded from a HuggingFace [tokenizer.json]\n file. Errors if the file cannot be read or has invalid format. *)\n\nval from_json : Jsont.json -> (t, string) result\n(** [from_json json] is a tokenizer deserialized from HuggingFace JSON format.\n Errors if [json] has a missing or unknown model type, or invalid parameters.\n*)\n\nval to_json : t -> Jsont.json\n(** [to_json t] is [t] serialized to HuggingFace JSON format. *)\n\nval save_pretrained : t -> path:string -> unit\n(** [save_pretrained t ~path] saves [t] to [path] in HuggingFace format. Creates\n [path/tokenizer.json].\n\n Raises [Sys_error] if [path] cannot be written. *)\n\n(** {1:fmt Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp] formats a tokenizer for inspection. Shows algorithm type, vocabulary\n size, and configured pipeline stages. *)\n"
},
{
"path": "packages/brot/lib/chars.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype t = unit\n\nlet create () = ()\n\nlet tokenize () text =\n if String.length text = 0 then []\n else\n let chars = ref [] in\n let offset = ref 0 in\n String.iter\n (fun c ->\n let char_str = String.make 1 c in\n let id = Char.code c in\n chars := (id, char_str, (!offset, !offset + 1)) :: !chars;\n incr offset)\n text;\n List.rev !chars\n\nlet token_to_id () token =\n if String.length token = 1 then Some (Char.code token.[0]) else None\n\nlet id_to_token () id =\n if id >= 0 && id <= 255 then Some (String.make 1 (Char.chr id)) else None\n\nlet get_vocab () = []\nlet get_vocab_size () = 256 (* All ASCII characters *)\nlet save () ~folder:_ () = []\n"
},
{
"path": "packages/brot/lib/chars.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Character-level tokenization model.\n\n {b Internal module.} Each byte maps to its ordinal value as token ID.\n Stateless: no vocabulary storage, no training. *)\n\ntype t\n(** The type for character-level models. *)\n\n(** {1:creation Creation} *)\n\nval create : unit -> t\n(** [create ()] is a character-level tokenizer. *)\n\n(** {1:tokenization Tokenization} *)\n\nval tokenize : t -> string -> (int * string * (int * int)) list\n(** [tokenize t s] is the tokenization of [s] as\n [(byte_value, char_string, (start, stop))] triples, one per byte. *)\n\n(** {1:vocabulary Vocabulary} *)\n\nval token_to_id : t -> string -> int option\n(** [token_to_id t s] is the byte value of [s] when [s] is a single byte. *)\n\nval id_to_token : t -> int -> string option\n(** [id_to_token t b] is the single-byte string for byte value [b]. *)\n\nval get_vocab : t -> (string * int) list\n(** [get_vocab t] is [[]] (no explicit vocabulary). *)\n\nval get_vocab_size : t -> int\n(** [get_vocab_size t] is [1114112] (all Unicode code points). *)\n\n(** {1:serialization Serialization} *)\n\nval save : t -> folder:string -> unit -> string list\n(** [save t ~folder ()] is [[]] (no files to write). *)\n"
},
{
"path": "packages/brot/lib/decoder.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype t =\n | BPE of { suffix : string }\n | Byte_level\n | Byte_fallback\n | Word_piece of { prefix : string; cleanup : bool }\n | Metaspace of { replacement : char; add_prefix_space : bool }\n | CTC of { pad_token : string; word_delimiter_token : string; cleanup : bool }\n | Sequence of t list\n | Replace of { pattern : string; replacement : string }\n | Strip of { left : bool; right : bool; content : char }\n | Fuse\n\n(* Errors *)\n\nlet strf = Printf.sprintf\nlet err_replace_missing_pattern = \"missing pattern in Replace decoder\"\n\nlet err_seq_missing_decoders =\n \"invalid Sequence decoder: missing decoders array\"\n\nlet err_unknown_type typ = strf \"unknown decoder type: %s\" typ\nlet err_expected_object = \"invalid decoder JSON: expected object\"\n\n(* Decoding *)\n\nlet whitespace_re = Re.compile (Re.rep1 (Re.char ' '))\n\n(* Literal string replacement without regex overhead. Returns [s] unchanged when\n [pattern] does not occur—no allocation on the fast path. *)\nlet replace_all ~pattern ~by s =\n let plen = String.length pattern in\n let slen = String.length s in\n if plen = 0 || plen > slen then s\n else\n let match_at i =\n let rec check j =\n j >= plen\n || String.unsafe_get s (i + j) = String.unsafe_get pattern j\n && check (j + 1)\n in\n check 0\n in\n let rec find_first i =\n if i > slen - plen then -1\n else if match_at i then i\n else find_first (i + 1)\n in\n let pos = find_first 0 in\n if pos < 0 then s\n else\n let buf = Buffer.create slen in\n Buffer.add_substring buf s 0 pos;\n Buffer.add_string buf by;\n let i = ref (pos + plen) in\n while !i <= slen - plen do\n if match_at !i then (\n Buffer.add_string buf by;\n i := !i + plen)\n else (\n Buffer.add_char buf (String.unsafe_get s !i);\n incr i)\n done;\n if !i < slen then Buffer.add_substring buf s !i (slen - !i);\n Buffer.contents buf\n\nlet decode_bpe ~suffix tokens =\n let suffix_len = String.length suffix in\n let strip token =\n if suffix_len > 0 && String.ends_with ~suffix token then\n String.sub token 0 (String.length token - suffix_len)\n else token\n in\n let rec loop acc = function\n | [] -> List.rev acc\n | [ token ] -> List.rev (strip token :: acc)\n | token :: rest -> loop (\" \" :: strip token :: acc) rest\n in\n loop [] tokens\n\nlet decode_byte_level tokens =\n let buf = Buffer.create 128 in\n List.iter\n (fun token -> Buffer.add_string buf (Pre_tokenizer.byte_level_decode token))\n tokens;\n Buffer.contents buf\n\nlet decode_byte_fallback tokens =\n let flush acc = function\n | [] -> acc\n | byte_acc ->\n let bytes = List.rev byte_acc in\n let s = Bytes.create (List.length bytes) in\n List.iteri (fun i b -> Bytes.unsafe_set s i (Char.chr b)) bytes;\n Bytes.unsafe_to_string s :: acc\n in\n let is_byte_token token =\n String.length token = 6\n && String.starts_with ~prefix:\"<0x\" token\n && String.ends_with ~suffix:\">\" token\n in\n let rec loop acc byte_acc = function\n | [] -> List.rev (flush acc byte_acc)\n | token :: rest when is_byte_token token -> (\n let hex = String.sub token 3 2 in\n match int_of_string_opt (\"0x\" ^ hex) with\n | Some b when b >= 0 && b <= 255 -> loop acc (b :: byte_acc) rest\n | _ -> loop (token :: flush acc byte_acc) [] rest)\n | token :: rest -> loop (token :: flush acc byte_acc) [] rest\n in\n loop [] [] tokens\n\nlet decode_wordpiece ~prefix ~cleanup tokens =\n let plen = String.length prefix in\n let buf = Buffer.create 128 in\n List.iteri\n (fun i token ->\n if i > 0 && String.starts_with ~prefix token then\n Buffer.add_substring buf token plen (String.length token - plen)\n else begin\n if i > 0 then Buffer.add_char buf ' ';\n Buffer.add_string buf token\n end)\n tokens;\n let s = Buffer.contents buf in\n if cleanup then String.trim (Re.replace_string whitespace_re ~by:\" \" s) else s\n\nlet decode_metaspace ~replacement ~add_prefix_space tokens =\n List.mapi\n (fun i token ->\n let s = String.map (fun c -> if c = replacement then ' ' else c) token in\n if add_prefix_space && i = 0 && String.length s > 0 && s.[0] = ' ' then\n String.sub s 1 (String.length s - 1)\n else s)\n tokens\n\nlet decode_ctc ~pad_token ~word_delimiter_token ~cleanup tokens =\n let rec dedup acc = function\n | [] -> List.rev acc\n | [ x ] -> List.rev (x :: acc)\n | x :: (y :: _ as rest) ->\n if String.equal x y then dedup acc rest else dedup (x :: acc) rest\n in\n let re =\n if cleanup then Some (Re.compile (Re.str word_delimiter_token)) else None\n in\n dedup [] tokens\n |> List.filter_map (fun token ->\n if String.equal token pad_token then None\n else\n let s =\n match re with\n | Some re -> Re.replace_string re ~by:\" \" token\n | None -> token\n in\n if String.equal s \"\" then None else Some s)\n\nlet decode_replace ~pattern ~replacement tokens =\n [ replace_all ~pattern ~by:replacement (String.concat \"\" tokens) ]\n\nlet strip_token ~left ~right content token =\n let len = String.length token in\n let start =\n if left then\n let rec find i =\n if i < len && Char.equal token.[i] content then find (i + 1) else i\n in\n find 0\n else 0\n in\n let stop =\n if right then\n let rec find i =\n if i >= 0 && Char.equal token.[i] content then find (i - 1) else i + 1\n in\n find (len - 1)\n else len\n in\n if start < stop then String.sub token start (stop - start) else \"\"\n\nlet rec decode_chain decoder tokens =\n match decoder with\n | BPE { suffix } -> decode_bpe ~suffix tokens\n | Byte_level -> [ decode_byte_level tokens ]\n | Byte_fallback -> decode_byte_fallback tokens\n | Word_piece { prefix; cleanup } ->\n [ decode_wordpiece ~prefix ~cleanup tokens ]\n | Metaspace { replacement; add_prefix_space } ->\n decode_metaspace ~replacement ~add_prefix_space tokens\n | CTC { pad_token; word_delimiter_token; cleanup } ->\n decode_ctc ~pad_token ~word_delimiter_token ~cleanup tokens\n | Replace { pattern; replacement } ->\n decode_replace ~pattern ~replacement tokens\n | Strip { left; right; content } ->\n [ strip_token ~left ~right content (String.concat \"\" tokens) ]\n | Fuse -> [ String.concat \"\" tokens ]\n | Sequence decoders ->\n List.fold_left (fun toks dec -> decode_chain dec toks) tokens decoders\n\nlet decode decoder tokens = String.concat \"\" (decode_chain decoder tokens)\n\n(* Constructors *)\n\nlet bpe ?(suffix = \"\") () = BPE { suffix }\nlet byte_level () = Byte_level\nlet byte_fallback () = Byte_fallback\n\nlet wordpiece ?(prefix = \"##\") ?(cleanup = true) () =\n Word_piece { prefix; cleanup }\n\nlet metaspace ?(replacement = '_') ?(add_prefix_space = true) () =\n Metaspace { replacement; add_prefix_space }\n\nlet ctc ?(pad_token = \"\") ?(word_delimiter_token = \"|\") ?(cleanup = true)\n () =\n CTC { pad_token; word_delimiter_token; cleanup }\n\nlet sequence decoders = Sequence decoders\nlet replace ~pattern ~by () = Replace { pattern; replacement = by }\n\nlet strip ?(left = false) ?(right = false) ?(content = ' ') () =\n Strip { left; right; content }\n\nlet fuse () = Fuse\n\n(* Formatting *)\n\nlet rec pp ppf = function\n | BPE { suffix } ->\n if suffix <> \"\" then Format.fprintf ppf \"bpe ~suffix:%S\" suffix\n else Format.fprintf ppf \"bpe\"\n | Byte_level -> Format.fprintf ppf \"byte_level\"\n | Byte_fallback -> Format.fprintf ppf \"byte_fallback\"\n | Word_piece { prefix; cleanup } ->\n Format.fprintf ppf \"wordpiece ~prefix:%S ~cleanup:%b\" prefix cleanup\n | Metaspace { replacement; add_prefix_space } ->\n Format.fprintf ppf \"metaspace ~replacement:%C ~add_prefix_space:%b\"\n replacement add_prefix_space\n | CTC { pad_token; word_delimiter_token; cleanup } ->\n Format.fprintf ppf\n \"ctc ~pad_token:%S ~word_delimiter_token:%S ~cleanup:%b\" pad_token\n word_delimiter_token cleanup\n | Replace { pattern; replacement } ->\n Format.fprintf ppf \"replace ~pattern:%S ~by:%S\" pattern replacement\n | Strip { left; right; content } ->\n Format.fprintf ppf \"strip ~left:%b ~right:%b ~content:%C\" left right\n content\n | Fuse -> Format.fprintf ppf \"fuse\"\n | Sequence decoders ->\n Format.fprintf ppf \"@[sequence [%a]@]\"\n (Format.pp_print_list\n ~pp_sep:(fun ppf () -> Format.fprintf ppf \";@ \")\n pp)\n decoders\n\n(* Serialization *)\n\nlet json_obj pairs =\n Jsont.Json.object' (List.map (fun (k, v) -> (Jsont.Json.name k, v)) pairs)\n\nlet rec to_json = function\n | BPE { suffix } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"BPEDecoder\");\n (\"suffix\", Jsont.Json.string suffix);\n ]\n | Byte_level -> json_obj [ (\"type\", Jsont.Json.string \"Byte_level\") ]\n | Byte_fallback -> json_obj [ (\"type\", Jsont.Json.string \"Byte_fallback\") ]\n | Word_piece { prefix; cleanup } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Word_piece\");\n (\"prefix\", Jsont.Json.string prefix);\n (\"cleanup\", Jsont.Json.bool cleanup);\n ]\n | Metaspace { replacement; add_prefix_space } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Metaspace\");\n (\"replacement\", Jsont.Json.string (String.make 1 replacement));\n (\"add_prefix_space\", Jsont.Json.bool add_prefix_space);\n ]\n | CTC { pad_token; word_delimiter_token; cleanup } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"CTC\");\n (\"pad_token\", Jsont.Json.string pad_token);\n (\"word_delimiter_token\", Jsont.Json.string word_delimiter_token);\n (\"cleanup\", Jsont.Json.bool cleanup);\n ]\n | Replace { pattern; replacement } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Replace\");\n (\"pattern\", Jsont.Json.string pattern);\n (\"content\", Jsont.Json.string replacement);\n ]\n | Strip { left; right; content } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Strip\");\n (\"strip_left\", Jsont.Json.bool left);\n (\"strip_right\", Jsont.Json.bool right);\n (\"content\", Jsont.Json.string (String.make 1 content));\n ]\n | Fuse -> json_obj [ (\"type\", Jsont.Json.string \"Fuse\") ]\n | Sequence decoders ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Sequence\");\n (\"decoders\", Jsont.Json.list (List.map to_json decoders));\n ]\n\nlet find_field fields name = Option.map snd (Jsont.Json.find_mem name fields)\n\nlet string_field fields name ~default =\n match find_field fields name with\n | Some (Jsont.String (s, _)) -> s\n | _ -> default\n\nlet bool_field fields name ~default =\n match find_field fields name with\n | Some (Jsont.Bool (b, _)) -> b\n | _ -> default\n\nlet char_field fields name ~default =\n match find_field fields name with\n | Some (Jsont.String (s, _)) when String.length s > 0 -> s.[0]\n | _ -> default\n\nlet rec of_json = function\n | Jsont.Object (fields, _) -> (\n let ( let* ) = Result.bind in\n match find_field fields \"type\" with\n | Some (Jsont.String (\"BPEDecoder\", _)) ->\n Ok (BPE { suffix = string_field fields \"suffix\" ~default:\"\" })\n | Some (Jsont.String ((\"Byte_level\" | \"ByteLevel\"), _)) -> Ok Byte_level\n | Some (Jsont.String ((\"Byte_fallback\" | \"ByteFallback\"), _)) ->\n Ok Byte_fallback\n | Some (Jsont.String ((\"Word_piece\" | \"WordPiece\"), _)) ->\n Ok\n (Word_piece\n {\n prefix = string_field fields \"prefix\" ~default:\"##\";\n cleanup = bool_field fields \"cleanup\" ~default:true;\n })\n | Some (Jsont.String (\"Metaspace\", _)) ->\n Ok\n (Metaspace\n {\n replacement = char_field fields \"replacement\" ~default:'_';\n add_prefix_space =\n bool_field fields \"add_prefix_space\" ~default:true;\n })\n | Some (Jsont.String (\"CTC\", _)) ->\n Ok\n (CTC\n {\n pad_token = string_field fields \"pad_token\" ~default:\"\";\n word_delimiter_token =\n string_field fields \"word_delimiter_token\" ~default:\"|\";\n cleanup = bool_field fields \"cleanup\" ~default:true;\n })\n | Some (Jsont.String (\"Replace\", _)) ->\n let* pattern =\n match find_field fields \"pattern\" with\n | Some (Jsont.String (s, _)) -> Ok s\n | Some (Jsont.Object (pattern_fields, _)) -> (\n match Jsont.Json.find_mem \"String\" pattern_fields with\n | Some (_, Jsont.String (p, _)) -> Ok p\n | _ -> Error err_replace_missing_pattern)\n | _ -> Error err_replace_missing_pattern\n in\n Ok\n (Replace\n {\n pattern;\n replacement = string_field fields \"content\" ~default:\"\";\n })\n | Some (Jsont.String (\"Strip\", _)) ->\n Ok\n (Strip\n {\n left = bool_field fields \"strip_left\" ~default:false;\n right = bool_field fields \"strip_right\" ~default:false;\n content = char_field fields \"content\" ~default:' ';\n })\n | Some (Jsont.String (\"Fuse\", _)) -> Ok Fuse\n | Some (Jsont.String (\"Sequence\", _)) -> (\n match find_field fields \"decoders\" with\n | Some (Jsont.Array (decs, _)) ->\n let* decoders =\n List.fold_left\n (fun acc j ->\n let* acc = acc in\n let* d = of_json j in\n Ok (d :: acc))\n (Ok []) decs\n in\n Ok (Sequence (List.rev decoders))\n | _ -> Error err_seq_missing_decoders)\n | Some (Jsont.String (typ, _)) -> Error (err_unknown_type typ)\n | _ -> Error \"missing or invalid decoder type field\")\n | _ -> Error err_expected_object\n"
},
{
"path": "packages/brot/lib/decoder.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Decoding tokens back to text.\n\n Decoders convert token strings back into natural text by reversing\n encoding-specific transformations (prefix/suffix removal, byte-level\n decoding, whitespace normalization, etc.).\n\n Decoders operate on token {e strings}, not IDs. Convert IDs to strings via\n vocabulary first, then apply {!decode}.\n\n Some decoders transform each token independently ({e per-token}: {!bpe},\n {!metaspace}, {!replace}, {!strip}, {!byte_fallback}), while others collapse\n the entire token list into a single result ({e collapsing}: {!byte_level},\n {!wordpiece}, {!fuse}). This distinction matters when composing decoders\n with {!sequence}. *)\n\ntype t\n(** The type for decoders. *)\n\n(** {1:constructors Constructors} *)\n\nval bpe : ?suffix:string -> unit -> t\n(** [bpe ~suffix ()] is a per-token decoder for BPE-encoded tokens. Strips\n [suffix] from end-of-word tokens and inserts spaces between words. [suffix]\n defaults to [\"\"]. *)\n\nval byte_level : unit -> t\n(** [byte_level ()] is a collapsing decoder that reverses GPT-2 style\n byte-to-Unicode encoding back to original bytes. *)\n\nval byte_fallback : unit -> t\n(** [byte_fallback ()] is a per-token decoder for byte fallback tokens. Converts\n hex byte tokens (e.g. [\"<0x41>\"]) back to their byte values, accumulating\n consecutive byte tokens into strings. Non-byte tokens pass through\n unchanged. *)\n\nval wordpiece : ?prefix:string -> ?cleanup:bool -> unit -> t\n(** [wordpiece ~prefix ~cleanup ()] is a collapsing decoder for WordPiece\n tokens. Strips continuation [prefix] (default [\"##\"]) from non-initial\n subwords and joins tokens into words. When [cleanup] is [true] (default),\n normalizes whitespace in the result. *)\n\nval metaspace : ?replacement:char -> ?add_prefix_space:bool -> unit -> t\n(** [metaspace ~replacement ~add_prefix_space ()] is a per-token decoder that\n converts metaspace markers back to regular spaces. [replacement] defaults to\n ['_']. When [add_prefix_space] is [true] (default), the leading replacement\n character on the first token is stripped. *)\n\nval ctc :\n ?pad_token:string ->\n ?word_delimiter_token:string ->\n ?cleanup:bool ->\n unit ->\n t\n(** [ctc ~pad_token ~word_delimiter_token ~cleanup ()] is a per-token decoder\n for\n {{:https://distill.pub/2017/ctc/}CTC (Connectionist Temporal\n Classification)} output. Deduplicates consecutive tokens, removes\n [pad_token] (default [\"\"]), and when [cleanup] is [true] (default),\n replaces [word_delimiter_token] (default [\"|\"]) with spaces. *)\n\nval sequence : t list -> t\n(** [sequence decoders] chains [decoders] left-to-right. Each decoder's output\n token list feeds into the next. *)\n\nval replace : pattern:string -> by:string -> unit -> t\n(** [replace ~pattern ~by ()] is a collapsing decoder that joins the token list,\n replaces all literal occurrences of [pattern] with [by] in the result, and\n returns a single-element list. *)\n\nval strip : ?left:bool -> ?right:bool -> ?content:char -> unit -> t\n(** [strip ~left ~right ~content ()] is a collapsing decoder that joins the\n token list and removes leading (when [left] is [true]) and/or trailing (when\n [right] is [true]) occurrences of [content] from the result. [left] and\n [right] default to [false]; [content] defaults to [' ']. *)\n\nval fuse : unit -> t\n(** [fuse ()] is a collapsing decoder that concatenates all tokens into a single\n string with no delimiter. *)\n\n(** {1:ops Operations} *)\n\nval decode : t -> string list -> string\n(** [decode decoder tokens] applies [decoder] to [tokens] and returns the\n decoded text. *)\n\n(** {1:fmt Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp ppf decoder] formats [decoder] for debugging. *)\n\n(** {1:serialization Serialization} *)\n\nval to_json : t -> Jsont.json\n(** [to_json decoder] serializes [decoder] to HuggingFace JSON format. *)\n\nval of_json : Jsont.json -> (t, string) result\n(** [of_json json] is a decoder from HuggingFace JSON format. Errors if [json]\n is not an object, has a missing or unknown [\"type\"] field, or has invalid\n parameters. *)\n"
},
{
"path": "packages/brot/lib/dune",
"content": "(library\n (name brot)\n (public_name brot)\n (private_modules bpe wordpiece word_level unigram chars)\n (libraries re jsont jsont.bytesrw uucp uunf))\n"
},
{
"path": "packages/brot/lib/encoding.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype t = {\n ids : int array;\n type_ids : int array;\n tokens : string array;\n words : int option array;\n offsets : (int * int) array;\n special_tokens_mask : int array;\n attention_mask : int array;\n mutable overflowing : t list;\n sequence_ranges : (int, int * int) Hashtbl.t;\n}\n\n(* Constructors *)\n\nlet empty_ranges : (int, int * int) Hashtbl.t = Hashtbl.create 0\n\nlet empty =\n {\n ids = [||];\n type_ids = [||];\n tokens = [||];\n words = [||];\n offsets = [||];\n special_tokens_mask = [||];\n attention_mask = [||];\n overflowing = [];\n sequence_ranges = empty_ranges;\n }\n\nlet create ~ids ~type_ids ~tokens ~words ~offsets ~special_tokens_mask\n ~attention_mask ?(overflowing = []) () =\n {\n ids;\n type_ids;\n tokens;\n words;\n offsets;\n special_tokens_mask;\n attention_mask;\n overflowing;\n sequence_ranges = empty_ranges;\n }\n\nlet token ~id ~token ~offset ~type_id ~special =\n {\n ids = [| id |];\n type_ids = [| type_id |];\n tokens = [| token |];\n words = [| None |];\n offsets = [| offset |];\n special_tokens_mask = [| (if special then 1 else 0) |];\n attention_mask = [| 1 |];\n overflowing = [];\n sequence_ranges = empty_ranges;\n }\n\nlet from_tokens tokens ~type_id =\n let n = List.length tokens in\n let ids = Array.make n 0 in\n let token_strs = Array.make n \"\" in\n let offsets = Array.make n (0, 0) in\n List.iteri\n (fun i (id, tok, off) ->\n ids.(i) <- id;\n token_strs.(i) <- tok;\n offsets.(i) <- off)\n tokens;\n {\n ids;\n tokens = token_strs;\n offsets;\n words = Array.make n None;\n type_ids = Array.make n type_id;\n attention_mask = Array.make n 1;\n special_tokens_mask = Array.make n 0;\n overflowing = [];\n sequence_ranges = empty_ranges;\n }\n\nlet concat a b =\n {\n ids = Array.append a.ids b.ids;\n type_ids = Array.append a.type_ids b.type_ids;\n tokens = Array.append a.tokens b.tokens;\n words = Array.append a.words b.words;\n offsets = Array.append a.offsets b.offsets;\n special_tokens_mask =\n Array.append a.special_tokens_mask b.special_tokens_mask;\n attention_mask = Array.append a.attention_mask b.attention_mask;\n overflowing = a.overflowing;\n sequence_ranges = a.sequence_ranges;\n }\n\nlet concat_list encodings =\n match encodings with\n | [] -> empty\n | [ single ] -> single\n | first :: _ ->\n let total =\n List.fold_left (fun acc t -> acc + Array.length t.ids) 0 encodings\n in\n let ids = Array.make total 0 in\n let type_ids = Array.make total 0 in\n let tokens = Array.make total \"\" in\n let words = Array.make total None in\n let offsets = Array.make total (0, 0) in\n let special_tokens_mask = Array.make total 0 in\n let attention_mask = Array.make total 0 in\n let pos = ref 0 in\n List.iter\n (fun t ->\n let n = Array.length t.ids in\n Array.blit t.ids 0 ids !pos n;\n Array.blit t.type_ids 0 type_ids !pos n;\n Array.blit t.tokens 0 tokens !pos n;\n Array.blit t.words 0 words !pos n;\n Array.blit t.offsets 0 offsets !pos n;\n Array.blit t.special_tokens_mask 0 special_tokens_mask !pos n;\n Array.blit t.attention_mask 0 attention_mask !pos n;\n pos := !pos + n)\n encodings;\n {\n ids;\n type_ids;\n tokens;\n words;\n offsets;\n special_tokens_mask;\n attention_mask;\n overflowing = first.overflowing;\n sequence_ranges = first.sequence_ranges;\n }\n\n(* Accessors *)\n\nlet is_empty t = Array.length t.ids = 0\nlet length t = Array.length t.ids\nlet ids t = t.ids\nlet type_ids t = t.type_ids\nlet tokens t = t.tokens\nlet word_ids t = t.words\nlet offsets t = t.offsets\nlet special_tokens_mask t = t.special_tokens_mask\nlet attention_mask t = t.attention_mask\nlet overflowing t = t.overflowing\n\n(* Truncation *)\n\nlet slice t start len =\n {\n ids = Array.sub t.ids start len;\n type_ids = Array.sub t.type_ids start len;\n tokens = Array.sub t.tokens start len;\n words = Array.sub t.words start len;\n offsets = Array.sub t.offsets start len;\n special_tokens_mask = Array.sub t.special_tokens_mask start len;\n attention_mask = Array.sub t.attention_mask start len;\n overflowing = [];\n sequence_ranges = empty_ranges;\n }\n\nlet truncate t ~max_length ~stride ~direction =\n let encoding_len = length t in\n if max_length >= encoding_len then t\n else if max_length = 0 then { empty with overflowing = [ t ] }\n else begin\n assert (stride < max_length);\n let step = max_length - stride in\n let ranges =\n match direction with\n | `Right ->\n let rec loop start acc =\n if start >= encoding_len then List.rev acc\n else\n let stop = min (start + max_length) encoding_len in\n loop (start + step) ((start, stop) :: acc)\n in\n loop 0 []\n | `Left ->\n let rec loop stop acc =\n if stop <= 0 then acc\n else\n let start = max 0 (stop - max_length) in\n loop (stop - step) ((start, stop) :: acc)\n in\n loop encoding_len []\n in\n match ranges with\n | [] -> empty\n | (start, stop) :: rest ->\n let enc = slice t start (stop - start) in\n enc.overflowing <-\n List.map (fun (start, stop) -> slice t start (stop - start)) rest;\n enc\n end\n\n(* Pad *)\n\nlet pad_array src n fill direction =\n let src_len = Array.length src in\n let dst = Array.make (src_len + n) fill in\n let off = match direction with `Left -> n | `Right -> 0 in\n Array.blit src 0 dst off src_len;\n dst\n\nlet rec pad t ~target_length ~pad_id ~pad_type_id ~pad_token ~direction =\n let overflowing =\n List.map\n (fun e -> pad e ~target_length ~pad_id ~pad_type_id ~pad_token ~direction)\n t.overflowing\n in\n let current_len = length t in\n if current_len >= target_length then { t with overflowing }\n else\n let n = target_length - current_len in\n let pad_a arr fill = pad_array arr n fill direction in\n let sequence_ranges =\n match direction with\n | `Right -> t.sequence_ranges\n | `Left ->\n if Hashtbl.length t.sequence_ranges = 0 then empty_ranges\n else begin\n let tbl = Hashtbl.create (Hashtbl.length t.sequence_ranges) in\n Hashtbl.iter\n (fun seq_id (start, stop) ->\n Hashtbl.add tbl seq_id (start + n, stop + n))\n t.sequence_ranges;\n tbl\n end\n in\n {\n ids = pad_a t.ids pad_id;\n type_ids = pad_a t.type_ids pad_type_id;\n tokens = pad_a t.tokens pad_token;\n words = pad_a t.words None;\n offsets = pad_a t.offsets (0, 0);\n special_tokens_mask = pad_a t.special_tokens_mask 1;\n attention_mask = pad_a t.attention_mask 0;\n overflowing;\n sequence_ranges;\n }\n"
},
{
"path": "packages/brot/lib/encoding.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Tokenization encodings.\n\n An encoding bundles token IDs for model input with alignment metadata: byte\n offsets, word indices, segment type IDs, attention masks, and special-token\n flags.\n\n Encodings are produced by {!Brot.encode} and post-processed with\n {!val-truncate} and {!val-pad}. All parallel arrays ({!val-ids},\n {!val-type_ids}, {!val-tokens}, {!val-word_ids}, {!val-offsets},\n {!val-special_tokens_mask}, {!val-attention_mask}) share the same length,\n equal to {!val-length}. *)\n\ntype t\n(** The type for tokenization encodings. *)\n\n(** {1:construct Construction} *)\n\nval empty : t\n(** [empty] is the encoding with no tokens. *)\n\nval create :\n ids:int array ->\n type_ids:int array ->\n tokens:string array ->\n words:int option array ->\n offsets:(int * int) array ->\n special_tokens_mask:int array ->\n attention_mask:int array ->\n ?overflowing:t list ->\n unit ->\n t\n(** [create ~ids ~type_ids ~tokens ~words ~offsets ~special_tokens_mask\n ~attention_mask ()] is an encoding from the given arrays.\n\n All arrays must have the same length; no validation is performed.\n [overflowing] defaults to [[]]. *)\n\nval token :\n id:int -> token:string -> offset:int * int -> type_id:int -> special:bool -> t\n(** [token ~id ~token ~offset ~type_id ~special] is a single-token encoding.\n When [special] is [true], {!val-special_tokens_mask} is [1] and\n {!val-word_ids} is [None]; otherwise {!val-special_tokens_mask} is [0].\n {!val-attention_mask} is always [1]. *)\n\nval from_tokens : (int * string * (int * int)) list -> type_id:int -> t\n(** [from_tokens tokens ~type_id] is an encoding from a list of\n [(id, token_string, (start, end_offset))] triples. Every token gets the\n given [type_id], {!val-attention_mask} [1], {!val-special_tokens_mask} [0]\n and {!val-word_ids} [None]. *)\n\nval concat : t -> t -> t\n(** [concat a b] is the encoding with [a]'s tokens followed by [b]'s.\n {!val-overflowing} and sequence ranges are taken from [a]. *)\n\nval concat_list : t list -> t\n(** [concat_list encs] is the concatenation of [encs] in order.\n {!val-overflowing} and sequence ranges are taken from the first element.\n Allocates once rather than creating intermediate arrays per pair. *)\n\n(** {1:access Accessors} *)\n\nval ids : t -> int array\n(** [ids enc] is the token ID array. *)\n\nval type_ids : t -> int array\n(** [type_ids enc] is the segment ID array. Typically [0] for the first sequence\n and [1] for the second in sentence-pair tasks. *)\n\nval tokens : t -> string array\n(** [tokens enc] is the string representation of each token. *)\n\nval word_ids : t -> int option array\n(** [word_ids enc] maps each token to its source word index, or [None] for\n special tokens. *)\n\nval offsets : t -> (int * int) array\n(** [offsets enc] is the [(start, end_)] byte offset spans into the original\n text for each token. *)\n\nval special_tokens_mask : t -> int array\n(** [special_tokens_mask enc] is [1] for special tokens ([CLS], [SEP], padding)\n and [0] for content tokens. *)\n\nval attention_mask : t -> int array\n(** [attention_mask enc] is [1] for real tokens and [0] for padding tokens. *)\n\nval overflowing : t -> t list\n(** [overflowing enc] is the list of overflow encodings produced by\n {!val-truncate} when the input exceeds [max_length]. Each element is a\n sliding window over the excess tokens. *)\n\nval is_empty : t -> bool\n(** [is_empty enc] is [true] iff [enc] has no tokens. *)\n\nval length : t -> int\n(** [length enc] is the number of tokens in [enc]. *)\n\n(** {1:ops Operations} *)\n\nval truncate :\n t -> max_length:int -> stride:int -> direction:[ `Left | `Right ] -> t\n(** [truncate enc ~max_length ~stride ~direction] limits [enc] to at most\n [max_length] tokens.\n\n Excess tokens are split into sliding windows of size [max_length] with\n overlap [stride] and stored in {!val-overflowing}. If\n [length enc <= max_length], [enc] is returned unchanged.\n\n [stride] must be strictly less than [max_length]. When [max_length] is [0],\n all tokens move to {!val-overflowing} and {!val-empty} is returned. *)\n\nval pad :\n t ->\n target_length:int ->\n pad_id:int ->\n pad_type_id:int ->\n pad_token:string ->\n direction:[ `Left | `Right ] ->\n t\n(** [pad enc ~target_length ~pad_id ~pad_type_id ~pad_token ~direction] extends\n [enc] to exactly [target_length] tokens.\n\n Padding tokens have {!val-attention_mask} [0] and {!val-special_tokens_mask}\n [1]. If [length enc >= target_length], [enc] is returned unchanged. Padding\n is applied recursively to {!val-overflowing} encodings. When [direction] is\n [`Left], {!val-offsets} and sequence ranges are shifted accordingly. *)\n"
},
{
"path": "packages/brot/lib/normalizer.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Errors *)\n\nlet err_expected_object = \"expected JSON object\"\nlet err_missing_type = \"missing type field\"\nlet err_replace_invalid_pattern = \"invalid pattern\"\nlet err_replace_missing_pattern = \"missing pattern\"\nlet err_replace_missing_content = \"missing content\"\nlet err_prepend_missing = \"missing prepend field\"\nlet err_sequence_missing = \"missing normalizers\"\nlet strf = Printf.sprintf\n\n(* Type *)\n\ntype t =\n | Bert of {\n clean_text : bool;\n handle_chinese_chars : bool;\n strip_accents : bool option;\n lowercase : bool;\n }\n | Strip of { left : bool; right : bool }\n | Strip_accents\n | NFC\n | NFD\n | NFKC\n | NFKD\n | Lowercase\n | Replace of { pattern : string; replacement : string; compiled : Re.re }\n | Prepend of string\n | Byte_level of { add_prefix_space : bool; use_regex : bool }\n | Sequence of t list\n\n(* Unicode text transforms *)\n\nlet normalize_utf8 nf text =\n let len = String.length text in\n if len = 0 then text\n else\n let rec all_ascii i =\n i >= len\n || (Char.code (String.unsafe_get text i) < 0x80 && all_ascii (i + 1))\n in\n if all_ascii 0 then text else Uunf_string.normalize_utf_8 nf text\n\nlet case_fold text =\n let len = String.length text in\n let rec needs_fold i =\n if i >= len then false\n else\n let byte = Char.code (String.unsafe_get text i) in\n if byte >= 0x41 && byte <= 0x5A then true\n else if byte >= 128 then true\n else needs_fold (i + 1)\n in\n if not (needs_fold 0) then text\n else\n let b = Buffer.create len in\n let i = ref 0 in\n while !i < len do\n let byte = Char.code (String.unsafe_get text !i) in\n if byte < 128 then (\n let c = if byte >= 0x41 && byte <= 0x5A then byte + 32 else byte in\n Buffer.add_char b (Char.unsafe_chr c);\n incr i)\n else\n let d = String.get_utf_8_uchar text !i in\n let n = Uchar.utf_decode_length d in\n (if Uchar.utf_decode_is_valid d then\n let u = Uchar.utf_decode_uchar d in\n match Uucp.Case.Fold.fold u with\n | `Self -> Buffer.add_utf_8_uchar b u\n | `Uchars us -> List.iter (fun u -> Buffer.add_utf_8_uchar b u) us);\n i := !i + n\n done;\n Buffer.contents b\n\nlet strip_accents_text text =\n let len = String.length text in\n let rec has_non_ascii i =\n if i >= len then false\n else if Char.code (String.unsafe_get text i) >= 128 then true\n else has_non_ascii (i + 1)\n in\n if not (has_non_ascii 0) then text\n else\n let b = Buffer.create len in\n let i = ref 0 in\n while !i < len do\n let byte = Char.code (String.unsafe_get text !i) in\n if byte < 128 then (\n Buffer.add_char b (Char.unsafe_chr byte);\n incr i)\n else\n let d = String.get_utf_8_uchar text !i in\n let n = Uchar.utf_decode_length d in\n (if Uchar.utf_decode_is_valid d then\n let u = Uchar.utf_decode_uchar d in\n match Uucp.Gc.general_category u with\n | `Mn | `Mc | `Me -> ()\n | _ -> Buffer.add_utf_8_uchar b u);\n i := !i + n\n done;\n Buffer.contents b\n\n(* UTF-8 helpers *)\n\n(* Returns (codepoint lsl 3) lor byte_length — zero allocation. *)\nlet[@inline] utf8_next s i =\n let d = String.get_utf_8_uchar s i in\n (Uchar.to_int (Uchar.utf_decode_uchar d) lsl 3) lor Uchar.utf_decode_length d\n\n(* Character classification *)\n\nlet[@inline] is_whitespace code =\n code = 0x09 || code = 0x0A || code = 0x0D || code = 0x20\n || Uucp.White.is_white_space (Uchar.of_int code)\n\nlet[@inline] is_control code =\n if code = 0x09 || code = 0x0A || code = 0x0D then false\n else\n match Uucp.Gc.general_category (Uchar.of_int code) with\n | `Cc | `Cf | `Cn | `Co -> true\n | _ -> false\n\nlet[@inline] is_chinese_char code =\n (code >= 0x4E00 && code <= 0x9FFF)\n || (code >= 0x3400 && code <= 0x4DBF)\n || (code >= 0x20000 && code <= 0x2A6DF)\n || (code >= 0x2A700 && code <= 0x2B73F)\n || (code >= 0x2B740 && code <= 0x2B81F)\n || (code >= 0x2B920 && code <= 0x2CEAF)\n || (code >= 0xF900 && code <= 0xFAFF)\n || (code >= 0x2F800 && code <= 0x2FA1F)\n\n(* Operations *)\n\nlet clean_text s =\n let len = String.length s in\n let buf = Buffer.create len in\n let i = ref 0 in\n while !i < len do\n let b0 = Char.code (String.unsafe_get s !i) in\n if b0 < 128 then begin\n if b0 = 9 || b0 = 10 || b0 = 13 || b0 = 32 then Buffer.add_char buf ' '\n else if b0 >= 33 && b0 < 127 then Buffer.add_char buf (Char.unsafe_chr b0);\n incr i\n end\n else begin\n let p = utf8_next s !i in\n let code = p lsr 3 and clen = p land 7 in\n if code <> 0xFFFD && not (is_control code) then\n if is_whitespace code then Buffer.add_char buf ' '\n else Buffer.add_substring buf s !i clen;\n i := !i + clen\n end\n done;\n Buffer.contents buf\n\nlet handle_chinese_chars s =\n let len = String.length s in\n let rec has_non_ascii i =\n i < len\n && (Char.code (String.unsafe_get s i) >= 128 || has_non_ascii (i + 1))\n in\n if not (has_non_ascii 0) then s\n else\n let buf = Buffer.create (len + (len / 4)) in\n let i = ref 0 in\n while !i < len do\n let b0 = Char.code (String.unsafe_get s !i) in\n if b0 < 128 then begin\n Buffer.add_char buf (Char.unsafe_chr b0);\n incr i\n end\n else begin\n let p = utf8_next s !i in\n let code = p lsr 3 and clen = p land 7 in\n if is_chinese_char code then (\n Buffer.add_char buf ' ';\n Buffer.add_substring buf s !i clen;\n Buffer.add_char buf ' ')\n else Buffer.add_substring buf s !i clen;\n i := !i + clen\n end\n done;\n Buffer.contents buf\n\nlet do_strip_accents s = strip_accents_text (normalize_utf8 `NFD s)\nlet do_lowercase s = case_fold s\n\nlet strip_whitespace s ~left ~right =\n let len = String.length s in\n let start =\n if left then\n let rec loop i =\n if i >= len then len\n else\n let p = utf8_next s i in\n let code = p lsr 3 and clen = p land 7 in\n if is_whitespace code then loop (i + clen) else i\n in\n loop 0\n else 0\n in\n let stop =\n if right then\n let rec loop i last =\n if i >= len then last\n else\n let p = utf8_next s i in\n let code = p lsr 3 and clen = p land 7 in\n let next = i + clen in\n if is_whitespace code then loop next last else loop next next\n in\n loop start start\n else len\n in\n if start = 0 && stop = len then s else String.sub s start (stop - start)\n\n(* Byte-level encoding *)\n\nlet byte_to_unicode =\n let is_direct b =\n (b >= 33 && b <= 126) || (b >= 161 && b <= 172) || b >= 174\n in\n let tbl = Array.make 256 0 in\n let n = ref 0 in\n for b = 0 to 255 do\n if is_direct b then tbl.(b) <- b\n else (\n tbl.(b) <- 256 + !n;\n incr n)\n done;\n tbl\n\nlet apply_byte_level s ~add_prefix_space ~use_regex:_ =\n let s =\n if add_prefix_space && String.length s > 0 then\n let code = utf8_next s 0 lsr 3 in\n if is_whitespace code then s else \" \" ^ s\n else s\n in\n let len = String.length s in\n let buf = Buffer.create (len * 2) in\n for i = 0 to len - 1 do\n let b = Char.code (String.unsafe_get s i) in\n Buffer.add_utf_8_uchar buf (Uchar.of_int byte_to_unicode.(b))\n done;\n Buffer.contents buf\n\n(* Constructors *)\n\nlet nfc = NFC\nlet nfd = NFD\nlet nfkc = NFKC\nlet nfkd = NFKD\nlet lowercase = Lowercase\nlet strip_accents = Strip_accents\nlet strip ?(left = true) ?(right = true) () = Strip { left; right }\n\nlet replace ~pattern ~replacement =\n Replace { pattern; replacement; compiled = Re.compile (Re.Pcre.re pattern) }\n\nlet prepend s = Prepend s\n\nlet byte_level ?(add_prefix_space = false) () =\n Byte_level { add_prefix_space; use_regex = false }\n\nlet bert ?(clean_text = true) ?(handle_chinese_chars = true)\n ?(strip_accents = None) ?(lowercase = true) () =\n Bert { clean_text; handle_chinese_chars; strip_accents; lowercase }\n\nlet sequence ns = Sequence ns\n\n(* Apply *)\n\nlet rec apply t s =\n match t with\n | NFC -> normalize_utf8 `NFC s\n | NFD -> normalize_utf8 `NFD s\n | NFKC -> normalize_utf8 `NFKC s\n | NFKD -> normalize_utf8 `NFKD s\n | Lowercase -> do_lowercase s\n | Strip_accents -> do_strip_accents s\n | Strip { left; right } -> strip_whitespace s ~left ~right\n | Replace { compiled; replacement; _ } ->\n Re.replace_string compiled ~by:replacement s\n | Prepend prefix -> if String.length s = 0 then s else prefix ^ s\n | Byte_level { add_prefix_space; use_regex } ->\n apply_byte_level s ~add_prefix_space ~use_regex\n | Bert\n {\n clean_text = ct;\n handle_chinese_chars = hcc;\n strip_accents = sa;\n lowercase = lc;\n } ->\n let s = if ct then clean_text s else s in\n let s = if hcc then handle_chinese_chars s else s in\n let do_strip = match sa with Some v -> v | None -> lc in\n let s = if do_strip then do_strip_accents s else s in\n if lc then do_lowercase s else s\n | Sequence ns -> List.fold_left (fun s n -> apply n s) s ns\n\n(* Formatting *)\n\nlet pp_bool_opt ppf = function\n | None -> Format.pp_print_string ppf \"None\"\n | Some b -> Format.fprintf ppf \"Some(%b)\" b\n\nlet rec pp ppf = function\n | NFC -> Format.pp_print_string ppf \"NFC\"\n | NFD -> Format.pp_print_string ppf \"NFD\"\n | NFKC -> Format.pp_print_string ppf \"NFKC\"\n | NFKD -> Format.pp_print_string ppf \"NFKD\"\n | Lowercase -> Format.pp_print_string ppf \"Lowercase\"\n | Strip_accents -> Format.pp_print_string ppf \"StripAccents\"\n | Strip { left; right } ->\n Format.fprintf ppf \"@[<1>Strip(left=%b,@ right=%b)@]\" left right\n | Replace { pattern; replacement; _ } ->\n Format.fprintf ppf \"@[<1>Replace(%S,@ %S)@]\" pattern replacement\n | Prepend s -> Format.fprintf ppf \"Prepend(%S)\" s\n | Byte_level { add_prefix_space; use_regex } ->\n Format.fprintf ppf \"@[<1>ByteLevel(add_prefix_space=%b,@ use_regex=%b)@]\"\n add_prefix_space use_regex\n | Bert { clean_text; handle_chinese_chars; strip_accents; lowercase } ->\n Format.fprintf ppf\n \"@[<1>Bert(clean_text=%b,@ handle_chinese_chars=%b,@ \\\n strip_accents=%a,@ lowercase=%b)@]\"\n clean_text handle_chinese_chars pp_bool_opt strip_accents lowercase\n | Sequence ns ->\n Format.fprintf ppf \"@[<1>Sequence[%a]@]\"\n (Format.pp_print_list\n ~pp_sep:(fun ppf () -> Format.fprintf ppf \",@ \")\n pp)\n ns\n\n(*---------------------------------------------------------------------------\n Serialization\n ---------------------------------------------------------------------------*)\n\nlet json_obj pairs =\n Jsont.Json.object' (List.map (fun (k, v) -> (Jsont.Json.name k, v)) pairs)\n\nlet typed name = json_obj [ (\"type\", Jsont.Json.string name) ]\nlet typed_with name pairs = json_obj ((\"type\", Jsont.Json.string name) :: pairs)\n\nlet rec to_json = function\n | Bert { clean_text; handle_chinese_chars; strip_accents; lowercase } ->\n typed_with \"Bert\"\n [\n (\"clean_text\", Jsont.Json.bool clean_text);\n (\"handle_chinese_chars\", Jsont.Json.bool handle_chinese_chars);\n ( \"strip_accents\",\n match strip_accents with\n | None -> Jsont.Json.null ()\n | Some b -> Jsont.Json.bool b );\n (\"lowercase\", Jsont.Json.bool lowercase);\n ]\n | Strip { left; right } ->\n typed_with \"Strip\"\n [\n (\"strip_left\", Jsont.Json.bool left);\n (\"strip_right\", Jsont.Json.bool right);\n ]\n | Strip_accents -> typed \"StripAccents\"\n | NFC -> typed \"NFC\"\n | NFD -> typed \"NFD\"\n | NFKC -> typed \"NFKC\"\n | NFKD -> typed \"NFKD\"\n | Lowercase -> typed \"Lowercase\"\n | Replace { pattern; replacement; _ } ->\n typed_with \"Replace\"\n [\n (\"pattern\", json_obj [ (\"String\", Jsont.Json.string pattern) ]);\n (\"content\", Jsont.Json.string replacement);\n ]\n | Prepend prefix ->\n typed_with \"Prepend\" [ (\"prepend\", Jsont.Json.string prefix) ]\n | Byte_level { add_prefix_space; use_regex } ->\n typed_with \"ByteLevel\"\n [\n (\"add_prefix_space\", Jsont.Json.bool add_prefix_space);\n (\"use_regex\", Jsont.Json.bool use_regex);\n ]\n | Sequence ns ->\n typed_with \"Sequence\"\n [ (\"normalizers\", Jsont.Json.list (List.map to_json ns)) ]\n\nlet rec of_json = function\n | Jsont.Object (fields, _) -> (\n let find name = Option.map snd (Jsont.Json.find_mem name fields) in\n let get_bool name default =\n match find name with Some (Jsont.Bool (b, _)) -> b | _ -> default\n in\n match find \"type\" with\n | Some (Jsont.String ((\"Bert\" | \"BertNormalizer\"), _)) ->\n let strip_accents =\n match find \"strip_accents\" with\n | Some (Jsont.Bool (b, _)) -> Some b\n | _ -> None\n in\n Ok\n (Bert\n {\n clean_text = get_bool \"clean_text\" true;\n handle_chinese_chars = get_bool \"handle_chinese_chars\" true;\n strip_accents;\n lowercase = get_bool \"lowercase\" true;\n })\n | Some (Jsont.String (\"Strip\", _)) ->\n Ok\n (Strip\n {\n left = get_bool \"strip_left\" false;\n right = get_bool \"strip_right\" true;\n })\n | Some (Jsont.String (\"StripAccents\", _)) -> Ok Strip_accents\n | Some (Jsont.String (\"NFC\", _)) -> Ok NFC\n | Some (Jsont.String (\"NFD\", _)) -> Ok NFD\n | Some (Jsont.String (\"NFKC\", _)) -> Ok NFKC\n | Some (Jsont.String (\"NFKD\", _)) -> Ok NFKD\n | Some (Jsont.String (\"Lowercase\", _)) -> Ok Lowercase\n | Some (Jsont.String (\"Replace\", _)) ->\n let pattern =\n match find \"pattern\" with\n | Some (Jsont.Object (pf, _)) -> (\n match Jsont.Json.find_mem \"String\" pf with\n | Some (_, Jsont.String (p, _)) -> Ok p\n | _ -> Error err_replace_invalid_pattern)\n | _ -> Error err_replace_missing_pattern\n in\n let replacement =\n match find \"content\" with\n | Some (Jsont.String (r, _)) -> Ok r\n | _ -> Error err_replace_missing_content\n in\n Result.bind pattern (fun p ->\n Result.map\n (fun r -> replace ~pattern:p ~replacement:r)\n replacement)\n | Some (Jsont.String (\"Prepend\", _)) -> (\n match find \"prepend\" with\n | Some (Jsont.String (p, _)) -> Ok (Prepend p)\n | _ -> Error err_prepend_missing)\n | Some (Jsont.String (\"ByteLevel\", _)) ->\n Ok\n (Byte_level\n {\n add_prefix_space = get_bool \"add_prefix_space\" false;\n use_regex = get_bool \"use_regex\" false;\n })\n | Some (Jsont.String (\"Sequence\", _)) -> (\n match find \"normalizers\" with\n | Some (Jsont.Array (l, _)) ->\n let rec build acc = function\n | [] -> Ok (Sequence (List.rev acc))\n | item :: rest ->\n Result.bind (of_json item) (fun n -> build (n :: acc) rest)\n in\n build [] l\n | _ -> Error err_sequence_missing)\n | Some (Jsont.String (other, _)) ->\n Error (strf \"Unknown normalizer type: %s\" other)\n | _ -> Error err_missing_type)\n | _ -> Error err_expected_object\n"
},
{
"path": "packages/brot/lib/normalizer.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Text normalization.\n\n Normalizers transform text before tokenization: lowercasing, accent removal,\n Unicode normalization, whitespace cleanup, and model-specific preprocessing.\n They are the first stage in the tokenization pipeline, applied before\n {!Pre_tokenizer} and vocabulary-based encoding.\n\n Compose normalizers with {!val-sequence}:\n {[\n let n =\n Normalizer.sequence\n [ Normalizer.nfd; Normalizer.strip_accents; Normalizer.lowercase ]\n in\n Normalizer.apply n \"Caf\\u{00E9}\"\n (* \"cafe\" *)\n ]}\n\n See {!Brot} for the full tokenization pipeline. *)\n\ntype t\n(** The type for normalizers. *)\n\n(** {1:normalizers Normalizers} *)\n\n(** {2:unicode Unicode normalization} *)\n\nval nfc : t\n(** [nfc] is Unicode NFC normalization (canonical composition). *)\n\nval nfd : t\n(** [nfd] is Unicode NFD normalization (canonical decomposition). *)\n\nval nfkc : t\n(** [nfkc] is Unicode NFKC normalization (compatibility composition). *)\n\nval nfkd : t\n(** [nfkd] is Unicode NFKD normalization (compatibility decomposition). *)\n\n(** {2:text Text transforms} *)\n\nval lowercase : t\n(** [lowercase] is Unicode case folding to lowercase. *)\n\nval strip_accents : t\n(** [strip_accents] removes combining marks after NFD decomposition. Applies\n {!val-nfd} before stripping. *)\n\nval strip : ?left:bool -> ?right:bool -> unit -> t\n(** [strip ?left ?right ()] is a normalizer that strips Unicode whitespace from\n text boundaries. [left] and [right] default to [true]. *)\n\nval replace : pattern:string -> replacement:string -> t\n(** [replace ~pattern ~replacement] is a normalizer that replaces all [pattern]\n matches with [replacement]. [pattern] is a PCRE regular expression, compiled\n once at construction time.\n\n Raises [Re.Pcre.Parse_error] if [pattern] is not valid PCRE. *)\n\nval prepend : string -> t\n(** [prepend s] is a normalizer that prepends [s] to non-empty text. Empty text\n is returned unchanged. *)\n\n(** {2:byte_level Byte-level encoding} *)\n\nval byte_level : ?add_prefix_space:bool -> unit -> t\n(** [byte_level ?add_prefix_space ()] is GPT-2 style byte-level encoding. Each\n byte is mapped to a printable Unicode codepoint using the GPT-2\n byte-to-unicode table.\n - [add_prefix_space] adds a space prefix when the text does not start with\n whitespace. Defaults to [false]. *)\n\n(** {2:model Model-specific} *)\n\nval bert :\n ?clean_text:bool ->\n ?handle_chinese_chars:bool ->\n ?strip_accents:bool option ->\n ?lowercase:bool ->\n unit ->\n t\n(** [bert ()] is a BERT normalizer.\n\n - [clean_text]: remove control characters and normalize whitespace. Default:\n [true].\n - [handle_chinese_chars]: pad CJK ideographs with spaces. Default: [true].\n - [strip_accents]: strip accents after NFD decomposition. When [None],\n accents are stripped iff [lowercase] is [true]. Default: [None].\n - [lowercase]: lowercase text via Unicode case folding. Default: [true]. *)\n\n(** {2:composition Composition} *)\n\nval sequence : t list -> t\n(** [sequence ns] is the composition of normalizers [ns], applied left to right.\n*)\n\n(** {1:applying Applying} *)\n\nval apply : t -> string -> string\n(** [apply n s] is [s] normalized by [n]. *)\n\n(** {1:formatting Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp ppf n] formats [n] for inspection. *)\n\n(** {1:serialization Serialization} *)\n\nval to_json : t -> Jsont.json\n(** [to_json n] is [n] serialized to HuggingFace-compatible JSON. *)\n\nval of_json : Jsont.json -> (t, string) result\n(** [of_json json] is a normalizer deserialized from HuggingFace JSON. Errors if\n [json] is not an object, has a missing or unknown [\"type\"] field, or has\n invalid parameters. *)\n"
},
{
"path": "packages/brot/lib/post_processor.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet strf = Printf.sprintf\nlet err_type_id tok = strf \"expected integer type id after ':' in '%s'\" tok\nlet err_piece tok = strf \"expected 'id' or 'id:type_id', got '%s'\" tok\nlet err_unknown_special tok = strf \"unknown special token '%s'\" tok\nlet err_mismatch tok = strf \"ids and tokens differ in length for '%s'\" tok\nlet err_expected what v = strf \"expected %s, got %s\" what v\nlet err_seq_id = \"sequence id must be \\\"A\\\", \\\"B\\\", 0 or 1\"\nlet err_type_id_field = \"expected number for 'type_id'\"\nlet err_missing_sequence = \"template references a sequence not provided\"\nlet err_pair_required = \"pair template required when two sequences are provided\"\nlet err_pair_must_ref_both = \"pair template must reference both $A and $B\"\nlet err_template_def = \"expected string, array or null for template\"\nlet err_unsupported_piece = \"expected Sequence or SpecialToken object\"\nlet err_special_missing_id = \"missing 'id' in SpecialToken\"\nlet err_special_missing_ids = \"missing 'ids' in special token\"\nlet err_special_entry = \"expected object for special token entry\"\n\n(* Types *)\n\ntype sequence_id = Sequence_a | Sequence_b\n\ntype template_piece =\n | Piece_sequence of { id : sequence_id; type_id : int }\n | Piece_special of { key : string; type_id : int }\n\ntype template = template_piece list\n\ntype special_token = {\n key : string;\n value_ids : int list;\n value_tokens : string list;\n}\n\ntype token = string * int\n\ntype t =\n | Bert of { sep : token; cls : token }\n | Roberta of {\n sep : token;\n cls : token;\n pad : token;\n trim_offsets : bool;\n add_prefix_space : bool;\n }\n | ByteLevel of { trim_offsets : bool }\n | Template of {\n single : template;\n pair : template option;\n special_tokens : special_token list;\n }\n | Sequence of t list\n\n(* Helpers *)\n\nlet special_token ~id ~token ~type_id =\n Encoding.token ~id ~token ~offset:(0, 0) ~type_id ~special:true\n\nlet with_type_id enc type_id =\n Encoding.create ~ids:(Encoding.ids enc)\n ~type_ids:(Array.make (Encoding.length enc) type_id)\n ~tokens:(Encoding.tokens enc) ~words:(Encoding.word_ids enc)\n ~offsets:(Encoding.offsets enc)\n ~special_tokens_mask:(Encoding.special_tokens_mask enc)\n ~attention_mask:(Encoding.attention_mask enc)\n ()\n\nlet is_ws = function\n | ' ' | '\\t' | '\\n' | '\\r' | '\\x0b' | '\\x0c' -> true\n | _ -> false\n\nlet build_special_lookup special_tokens =\n let tbl = Hashtbl.create (List.length special_tokens + 1) in\n List.iter (fun tok -> Hashtbl.replace tbl tok.key tok) special_tokens;\n tbl\n\nlet string_is_int s =\n let len = String.length s in\n let rec loop i =\n if i >= len then true\n else match s.[i] with '0' .. '9' -> loop (i + 1) | _ -> false\n in\n len > 0 && loop 0\n\nlet sequence_id_to_label = function Sequence_a -> \"A\" | Sequence_b -> \"B\"\nlet sequence_id_to_index = function Sequence_a -> 0 | Sequence_b -> 1\n\n(* JSON helpers *)\n\nlet json_obj pairs =\n Jsont.Json.object' (List.map (fun (k, v) -> (Jsont.Json.name k, v)) pairs)\n\nlet json_find name fields =\n match Jsont.Json.find_mem name fields with\n | Some (_, v) -> Some v\n | None -> None\n\nlet json_bool_field fields name ~default =\n match json_find name fields with\n | Some (Jsont.Bool (b, _)) -> b\n | _ -> default\n\nlet json_str_int_pair fields name ~default =\n match json_find name fields with\n | Some (Jsont.Array ([ Jsont.String (s, _); Jsont.Number (f, _) ], _)) ->\n (s, int_of_float f)\n | _ -> default\n\n(* Processors *)\n\nlet process_bert ~sep ~cls encodings ~add_special_tokens =\n if not add_special_tokens then encodings\n else\n let cls_str, cls_id = cls in\n let sep_str, sep_id = sep in\n let cls_tok tid = special_token ~id:cls_id ~token:cls_str ~type_id:tid in\n let sep_tok tid = special_token ~id:sep_id ~token:sep_str ~type_id:tid in\n match encodings with\n | [] -> []\n | [ encoding ] ->\n [\n Encoding.concat_list [ cls_tok 0; with_type_id encoding 0; sep_tok 0 ];\n ]\n | [ enc1; enc2 ] ->\n [\n Encoding.concat_list\n [\n cls_tok 0;\n with_type_id enc1 0;\n sep_tok 0;\n with_type_id enc2 1;\n sep_tok 1;\n ];\n ]\n | _ -> encodings\n\nlet process_roberta ~sep ~cls ~pad:_ ~trim_offsets:_ ~add_prefix_space:_\n encodings ~add_special_tokens =\n if not add_special_tokens then encodings\n else\n let cls_str, cls_id = cls in\n let sep_str, sep_id = sep in\n let cls_tok = special_token ~id:cls_id ~token:cls_str ~type_id:0 in\n let sep_tok = special_token ~id:sep_id ~token:sep_str ~type_id:0 in\n match encodings with\n | [] -> []\n | [ encoding ] ->\n [ Encoding.concat_list [ cls_tok; with_type_id encoding 0; sep_tok ] ]\n | [ enc1; enc2 ] ->\n [\n Encoding.concat_list\n [\n cls_tok;\n with_type_id enc1 0;\n sep_tok;\n sep_tok;\n with_type_id enc2 0;\n sep_tok;\n ];\n ]\n | _ -> encodings\n\nlet trim_offset enc_tokens idx (start, stop) =\n if start >= stop then (start, stop)\n else\n let token =\n if idx < Array.length enc_tokens then enc_tokens.(idx) else \"\"\n in\n let decoded = Pre_tokenizer.byte_level_decode token in\n let len = String.length decoded in\n let rec leading i =\n if i >= len then len else if is_ws decoded.[i] then leading (i + 1) else i\n in\n let rec trailing i =\n if i <= 0 then len\n else if is_ws decoded.[i - 1] then trailing (i - 1)\n else i\n in\n let lead = leading 0 in\n let trail = trailing len in\n let trimmed_lead = min (stop - start) lead in\n let trimmed_trail = min (stop - start - trimmed_lead) (len - trail) in\n let new_start = start + trimmed_lead in\n let new_stop = max new_start (stop - trimmed_trail) in\n (new_start, new_stop)\n\nlet process_byte_level ~trim_offsets encodings ~add_special_tokens:_ =\n if not trim_offsets then encodings\n else\n List.map\n (fun encoding ->\n let enc_tokens = Encoding.tokens encoding in\n let new_offsets =\n Array.mapi (trim_offset enc_tokens) (Encoding.offsets encoding)\n in\n Encoding.create ~ids:(Encoding.ids encoding)\n ~type_ids:(Encoding.type_ids encoding)\n ~tokens:enc_tokens\n ~words:(Encoding.word_ids encoding)\n ~offsets:new_offsets\n ~special_tokens_mask:(Encoding.special_tokens_mask encoding)\n ~attention_mask:(Encoding.attention_mask encoding)\n ~overflowing:(Encoding.overflowing encoding)\n ())\n encodings\n\n(* Template parsing *)\n\nlet split_template_string str =\n let len = String.length str in\n let rec skip_ws i =\n if i >= len then len\n else match str.[i] with ' ' | '\\t' -> skip_ws (i + 1) | _ -> i\n in\n let rec find_end i =\n if i >= len then len\n else match str.[i] with ' ' | '\\t' -> i | _ -> find_end (i + 1)\n in\n let rec loop i acc =\n let i = skip_ws i in\n if i >= len then List.rev acc\n else\n let j = find_end i in\n loop j (String.sub str i (j - i) :: acc)\n in\n loop 0 []\n\nlet parse_sequence_base base =\n let lower = String.lowercase_ascii base in\n if lower = \"$\" || lower = \"$a\" then Some (Sequence_a, 0)\n else if lower = \"$b\" then Some (Sequence_b, 0)\n else if String.length base > 0 && base.[0] = '$' then\n let rest = String.sub base 1 (String.length base - 1) in\n if string_is_int rest then Some (Sequence_a, int_of_string rest) else None\n else None\n\nlet parse_template_piece_from_string ~special_lookup token =\n let parts = String.split_on_char ':' token in\n let base, explicit_type =\n match parts with\n | [ id; type_part ] when string_is_int type_part ->\n (id, Some (int_of_string type_part))\n | [ _; _ ] -> invalid_arg (err_type_id token)\n | [ id ] -> (id, None)\n | _ -> invalid_arg (err_piece token)\n in\n match parse_sequence_base base with\n | Some (seq_id, default_type) ->\n let type_id = Option.value ~default:default_type explicit_type in\n Piece_sequence { id = seq_id; type_id }\n | None ->\n if Hashtbl.mem special_lookup base then\n let type_id = Option.value ~default:0 explicit_type in\n Piece_special { key = base; type_id }\n else invalid_arg (err_unknown_special token)\n\nlet parse_template_string ~special_lookup str =\n List.map\n (parse_template_piece_from_string ~special_lookup)\n (split_template_string str)\n\nlet parse_sequence_id_json fields =\n match json_find \"id\" fields with\n | Some (Jsont.String (s, _)) -> (\n match String.lowercase_ascii s with\n | \"a\" -> Sequence_a\n | \"b\" -> Sequence_b\n | _ -> invalid_arg err_seq_id)\n | Some (Jsont.Number (v, _)) -> (\n match int_of_float v with\n | 0 -> Sequence_a\n | 1 -> Sequence_b\n | _ -> invalid_arg err_seq_id)\n | None -> Sequence_a\n | _ -> invalid_arg err_seq_id\n\nlet json_type_id fields =\n match json_find \"type_id\" fields with\n | Some (Jsont.Number (v, _)) -> int_of_float v\n | None -> 0\n | _ -> invalid_arg err_type_id_field\n\nlet parse_template_piece_from_json ~special_lookup json =\n match json with\n | Jsont.Object (outer_fields, _) -> (\n match json_find \"Sequence\" outer_fields with\n | Some (Jsont.Object (fields, _)) ->\n let id = parse_sequence_id_json fields in\n let type_id = json_type_id fields in\n Piece_sequence { id; type_id }\n | _ -> (\n match json_find \"SpecialToken\" outer_fields with\n | Some (Jsont.Object (fields, _)) ->\n let key =\n match json_find \"id\" fields with\n | Some (Jsont.String (s, _)) -> s\n | _ -> invalid_arg err_special_missing_id\n in\n if not (Hashtbl.mem special_lookup key) then\n invalid_arg (err_unknown_special key);\n let type_id = json_type_id fields in\n Piece_special { key; type_id }\n | _ -> invalid_arg err_unsupported_piece))\n | _ -> invalid_arg err_unsupported_piece\n\nlet parse_template_definition ~special_lookup = function\n | Jsont.String (s, _) -> parse_template_string ~special_lookup s\n | Jsont.Array (l, _) ->\n List.map (parse_template_piece_from_json ~special_lookup) l\n | Jsont.Null _ -> []\n | _ -> invalid_arg err_template_def\n\n(* Template encoding *)\n\nlet build_encoding_from_pieces pieces source_encodings special_lookup =\n let ids_rev = ref [] in\n let type_ids_rev = ref [] in\n let tokens_rev = ref [] in\n let words_rev = ref [] in\n let offsets_rev = ref [] in\n let special_mask_rev = ref [] in\n let attention_rev = ref [] in\n let append ~id ~token ~word ~type_id ~offset ~special ~attention =\n ids_rev := id :: !ids_rev;\n type_ids_rev := type_id :: !type_ids_rev;\n tokens_rev := token :: !tokens_rev;\n words_rev := word :: !words_rev;\n offsets_rev := offset :: !offsets_rev;\n special_mask_rev := special :: !special_mask_rev;\n attention_rev := attention :: !attention_rev\n in\n let append_sequence seq_id type_id =\n let index = sequence_id_to_index seq_id in\n if index >= Array.length source_encodings then\n invalid_arg err_missing_sequence;\n let src = source_encodings.(index) in\n let src_ids = Encoding.ids src in\n let src_tokens = Encoding.tokens src in\n let src_words = Encoding.word_ids src in\n let src_offsets = Encoding.offsets src in\n let src_special = Encoding.special_tokens_mask src in\n let src_attention = Encoding.attention_mask src in\n let len = Array.length src_ids in\n for i = 0 to len - 1 do\n let token = if i < Array.length src_tokens then src_tokens.(i) else \"\" in\n let word = if i < Array.length src_words then src_words.(i) else None in\n let offset =\n if i < Array.length src_offsets then src_offsets.(i) else (0, 0)\n in\n let special =\n if i < Array.length src_special && src_special.(i) <> 0 then 1 else 0\n in\n let attention =\n if i < Array.length src_attention && src_attention.(i) <> 0 then 1\n else 0\n in\n append ~id:src_ids.(i) ~token ~word ~type_id ~offset ~special ~attention\n done\n in\n let append_special key type_id =\n match Hashtbl.find_opt special_lookup key with\n | None -> invalid_arg (err_unknown_special key)\n | Some special ->\n let rec loop ids tokens =\n match (ids, tokens) with\n | id :: rest_ids, token :: rest_tokens ->\n append ~id ~token ~word:None ~type_id ~offset:(0, 0) ~special:1\n ~attention:1;\n loop rest_ids rest_tokens\n | [], [] -> ()\n | _ -> invalid_arg (err_mismatch key)\n in\n loop special.value_ids special.value_tokens\n in\n List.iter\n (function\n | Piece_sequence { id; type_id } -> append_sequence id type_id\n | Piece_special { key; type_id } -> append_special key type_id)\n pieces;\n let to_array r = Array.of_list (List.rev !r) in\n Encoding.create ~ids:(to_array ids_rev) ~type_ids:(to_array type_ids_rev)\n ~tokens:(to_array tokens_rev) ~words:(to_array words_rev)\n ~offsets:(to_array offsets_rev)\n ~special_tokens_mask:(to_array special_mask_rev)\n ~attention_mask:(to_array attention_rev) ()\n\nlet process_template ~single ~pair ~special_tokens encodings ~add_special_tokens\n =\n if not add_special_tokens then encodings\n else\n let special_lookup = build_special_lookup special_tokens in\n let source = Array.of_list encodings in\n match Array.length source with\n | 0 -> []\n | 1 -> [ build_encoding_from_pieces single source special_lookup ]\n | 2 ->\n let pair =\n match pair with Some p -> p | None -> invalid_arg err_pair_required\n in\n [ build_encoding_from_pieces pair source special_lookup ]\n | _ -> encodings\n\n(* Processing *)\n\nlet rec process_list processor encodings ~add_special_tokens =\n match processor with\n | Bert { sep; cls } -> process_bert ~sep ~cls encodings ~add_special_tokens\n | Roberta { sep; cls; pad; trim_offsets; add_prefix_space } ->\n process_roberta ~sep ~cls ~pad ~trim_offsets ~add_prefix_space encodings\n ~add_special_tokens\n | ByteLevel { trim_offsets } ->\n process_byte_level ~trim_offsets encodings ~add_special_tokens\n | Template { single; pair; special_tokens } ->\n process_template ~single ~pair ~special_tokens encodings\n ~add_special_tokens\n | Sequence processors ->\n List.fold_left\n (fun encs proc -> process_list proc encs ~add_special_tokens)\n encodings processors\n\nlet process processor ?pair enc ~add_special_tokens =\n let encodings = match pair with None -> [ enc ] | Some p -> [ enc; p ] in\n match process_list processor encodings ~add_special_tokens with\n | [ r ] -> r\n | r :: _ -> r\n | [] -> enc\n\nlet rec added_tokens processor ~is_pair =\n match processor with\n | Bert _ -> if is_pair then 3 else 2\n | Roberta _ -> if is_pair then 4 else 2\n | ByteLevel _ -> 0\n | Template { single; pair; special_tokens } ->\n let lookup = build_special_lookup special_tokens in\n let count_special pieces =\n List.fold_left\n (fun acc piece ->\n match piece with\n | Piece_special { key; _ } -> (\n match Hashtbl.find_opt lookup key with\n | Some tok -> acc + List.length tok.value_ids\n | None -> acc)\n | _ -> acc)\n 0 pieces\n in\n if is_pair then\n match pair with\n | Some p -> count_special p\n | None -> count_special single\n else count_special single\n | Sequence processors ->\n List.fold_left\n (fun acc proc -> acc + added_tokens proc ~is_pair)\n 0 processors\n\n(* Constructors *)\n\nlet bert ~sep ~cls () = Bert { sep; cls }\n\nlet roberta ~sep ~cls ?(trim_offsets = true) ?(add_prefix_space = true) () =\n let pad = (\"\", 1) in\n Roberta { sep; cls; pad; trim_offsets; add_prefix_space }\n\nlet byte_level ?(trim_offsets = true) () = ByteLevel { trim_offsets }\n\nlet template ~single ?pair ?(special_tokens = []) () =\n let specials =\n List.map\n (fun (token, id) ->\n { key = token; value_ids = [ id ]; value_tokens = [ token ] })\n special_tokens\n in\n let lookup = build_special_lookup specials in\n let single = parse_template_string ~special_lookup:lookup single in\n let has_sequence pieces seq =\n List.exists\n (function Piece_sequence { id; _ } when id = seq -> true | _ -> false)\n pieces\n in\n let pair =\n match pair with\n | None -> None\n | Some p ->\n let tpl = parse_template_string ~special_lookup:lookup p in\n if not (has_sequence tpl Sequence_a && has_sequence tpl Sequence_b) then\n invalid_arg err_pair_must_ref_both;\n Some tpl\n in\n Template { single; pair; special_tokens = specials }\n\nlet sequence processors = Sequence processors\n\n(* Formatting *)\n\nlet rec pp ppf = function\n | Bert { sep = sep_s, _; cls = cls_s, _ } ->\n Format.fprintf ppf \"@[<2>Bert@ ~cls:%S@ ~sep:%S@]\" cls_s sep_s\n | Roberta { sep = sep_s, _; cls = cls_s, _; _ } ->\n Format.fprintf ppf \"@[<2>Roberta@ ~cls:%S@ ~sep:%S@]\" cls_s sep_s\n | ByteLevel { trim_offsets } ->\n Format.fprintf ppf \"@[<2>ByteLevel@ ~trim_offsets:%b@]\" trim_offsets\n | Template _ -> Format.fprintf ppf \"Template\"\n | Sequence processors ->\n Format.fprintf ppf \"@[<2>Sequence[@,%a]@]\"\n (Format.pp_print_list\n ~pp_sep:(fun ppf () -> Format.fprintf ppf \",@ \")\n pp)\n processors\n\n(* Serialization *)\n\nlet token_pair_to_json (s, id) =\n Jsont.Json.list [ Jsont.Json.string s; Jsont.Json.int id ]\n\nlet template_to_json pieces =\n let piece_json tag id type_id =\n json_obj\n [ (tag, json_obj [ (\"id\", id); (\"type_id\", Jsont.Json.int type_id) ]) ]\n in\n Jsont.Json.list\n (List.map\n (function\n | Piece_sequence { id; type_id } ->\n piece_json \"Sequence\"\n (Jsont.Json.string (sequence_id_to_label id))\n type_id\n | Piece_special { key; type_id } ->\n piece_json \"SpecialToken\" (Jsont.Json.string key) type_id)\n pieces)\n\nlet rec to_json = function\n | Bert { sep; cls } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"BertProcessing\");\n (\"sep\", token_pair_to_json sep);\n (\"cls\", token_pair_to_json cls);\n ]\n | Roberta { sep; cls; pad; trim_offsets; add_prefix_space } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"RobertaProcessing\");\n (\"sep\", token_pair_to_json sep);\n (\"cls\", token_pair_to_json cls);\n (\"pad\", token_pair_to_json pad);\n (\"trim_offsets\", Jsont.Json.bool trim_offsets);\n (\"add_prefix_space\", Jsont.Json.bool add_prefix_space);\n ]\n | ByteLevel { trim_offsets } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"ByteLevel\");\n (\"trim_offsets\", Jsont.Json.bool trim_offsets);\n ]\n | Template { single; pair; special_tokens } ->\n let pair_json =\n match pair with\n | None -> Jsont.Json.null ()\n | Some p -> template_to_json p\n in\n let special_token_json tok =\n let ids = Jsont.Json.list (List.map Jsont.Json.int tok.value_ids) in\n let tokens =\n Jsont.Json.list (List.map Jsont.Json.string tok.value_tokens)\n in\n ( Jsont.Json.name tok.key,\n json_obj\n [\n (\"id\", Jsont.Json.string tok.key); (\"ids\", ids); (\"tokens\", tokens);\n ] )\n in\n let special_json =\n Jsont.Json.object' (List.map special_token_json special_tokens)\n in\n json_obj\n [\n (\"type\", Jsont.Json.string \"TemplateProcessing\");\n (\"single\", template_to_json single);\n (\"pair\", pair_json);\n (\"special_tokens\", special_json);\n ]\n | Sequence processors ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Sequence\");\n (\"processors\", Jsont.Json.list (List.map to_json processors));\n ]\n\n(* Deserialization *)\n\nlet parse_special_token_json fields alias =\n let key =\n match json_find \"id\" fields with\n | Some (Jsont.String (s, _)) -> s\n | _ -> alias\n in\n let value_ids =\n match json_find \"ids\" fields with\n | Some (Jsont.Array (lst, _)) ->\n List.map\n (function\n | Jsont.Number (f, _) -> int_of_float f\n | v ->\n invalid_arg\n (err_expected \"number\" (Format.asprintf \"%a\" Jsont.pp_json v)))\n lst\n | _ -> invalid_arg err_special_missing_ids\n in\n let value_tokens =\n match json_find \"tokens\" fields with\n | Some (Jsont.Array (lst, _)) ->\n List.map\n (function\n | Jsont.String (s, _) -> s\n | v ->\n invalid_arg\n (err_expected \"string\" (Format.asprintf \"%a\" Jsont.pp_json v)))\n lst\n | _ -> [ key ]\n in\n if List.length value_ids <> List.length value_tokens then\n invalid_arg (err_mismatch key);\n { key; value_ids; value_tokens }\n\nlet parse_special_tokens_json fields =\n match json_find \"special_tokens\" fields with\n | Some (Jsont.Object (tokens, _)) ->\n List.map\n (fun ((alias, _), value) ->\n match value with\n | Jsont.Object (token_fields, _) ->\n parse_special_token_json token_fields alias\n | _ -> invalid_arg err_special_entry)\n tokens\n | Some v ->\n invalid_arg\n (err_expected \"object for 'special_tokens'\"\n (Format.asprintf \"%a\" Jsont.pp_json v))\n | None -> []\n\nlet rec of_json_exn json =\n match json with\n | Jsont.Object (fields, _) -> (\n match json_find \"type\" fields with\n | Some (Jsont.String (\"BertProcessing\", _)) ->\n let sep = json_str_int_pair fields \"sep\" ~default:(\"[SEP]\", 102) in\n let cls = json_str_int_pair fields \"cls\" ~default:(\"[CLS]\", 101) in\n Bert { sep; cls }\n | Some (Jsont.String (\"RobertaProcessing\", _)) ->\n let sep = json_str_int_pair fields \"sep\" ~default:(\"\", 2) in\n let cls = json_str_int_pair fields \"cls\" ~default:(\"\", 0) in\n let pad = json_str_int_pair fields \"pad\" ~default:(\"\", 1) in\n let trim_offsets =\n json_bool_field fields \"trim_offsets\" ~default:true\n in\n let add_prefix_space =\n json_bool_field fields \"add_prefix_space\" ~default:true\n in\n Roberta { sep; cls; pad; trim_offsets; add_prefix_space }\n | Some (Jsont.String (\"ByteLevel\", _)) ->\n let trim_offsets =\n json_bool_field fields \"trim_offsets\" ~default:true\n in\n ByteLevel { trim_offsets }\n | Some (Jsont.String (\"TemplateProcessing\", _)) ->\n let special_tokens = parse_special_tokens_json fields in\n let lookup = build_special_lookup special_tokens in\n let single =\n match json_find \"single\" fields with\n | Some json -> parse_template_definition ~special_lookup:lookup json\n | None -> parse_template_string ~special_lookup:lookup \"$A\"\n in\n let pair =\n match json_find \"pair\" fields with\n | Some (Jsont.Null _) | None -> None\n | Some json ->\n Some (parse_template_definition ~special_lookup:lookup json)\n in\n Template { single; pair; special_tokens }\n | Some (Jsont.String (\"Sequence\", _)) -> (\n match json_find \"processors\" fields with\n | Some (Jsont.Array (procs, _)) ->\n Sequence (List.map of_json_exn procs)\n | _ -> failwith \"expected array for 'processors'\")\n | _ -> failwith \"unsupported processor type\")\n | _ -> failwith \"expected JSON object\"\n\nlet of_json json =\n try Ok (of_json_exn json) with\n | Failure msg -> Error msg\n | Invalid_argument msg -> Error msg\n"
},
{
"path": "packages/brot/lib/post_processor.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Post-processing tokenization output with special tokens.\n\n Post-processors add special tokens and type IDs to tokenized sequences after\n core tokenization. They handle model-specific requirements like [[CLS]] and\n [[SEP]] for BERT, sentence pair formatting, and byte-level offset\n adjustments. *)\n\ntype t\n(** The type for post-processors. *)\n\ntype token = string * int\n(** A special token as [(text, id)]. *)\n\n(** {1:constructors Constructors} *)\n\nval bert : sep:token -> cls:token -> unit -> t\n(** [bert ~sep ~cls ()] is a BERT-style post-processor.\n\n Single: [[CLS] A [SEP]]. Pair: [[CLS] A [SEP] B [SEP]]. Type IDs: [0] for\n the first sequence, [1] for the second. *)\n\nval roberta :\n sep:token ->\n cls:token ->\n ?trim_offsets:bool ->\n ?add_prefix_space:bool ->\n unit ->\n t\n(** [roberta ~sep ~cls ()] is a RoBERTa-style post-processor.\n\n Single: [ A ]. Pair: [ A B ]. All type IDs are [0].\n\n [trim_offsets] defaults to [true]. [add_prefix_space] defaults to [true]. *)\n\nval byte_level : ?trim_offsets:bool -> unit -> t\n(** [byte_level ()] is a byte-level post-processor that adjusts character\n offsets for byte-level encoding.\n\n [trim_offsets] removes leading and trailing whitespace from offsets.\n Defaults to [true]. *)\n\nval template :\n single:string -> ?pair:string -> ?special_tokens:token list -> unit -> t\n(** [template ~single ()] is a template-based post-processor.\n\n Templates use [$A] and [$B] as sequence placeholders and literal special\n token names (e.g. [[CLS]]). Type IDs can be specified with a colon suffix:\n [$A:0], [[SEP]:1].\n\n [special_tokens] defaults to [[]]. *)\n\nval sequence : t list -> t\n(** [sequence processors] chains [processors] left-to-right. *)\n\n(** {1:processing Processing} *)\n\nval process :\n t -> ?pair:Encoding.t -> Encoding.t -> add_special_tokens:bool -> Encoding.t\n(** [process t enc ~add_special_tokens] adds special tokens and sets type IDs on\n [enc].\n\n When [~pair] is provided, both sequences are merged into a single encoding\n with appropriate type IDs. When [~add_special_tokens] is [false], special\n token insertion is skipped but byte-level offset trimming still applies. *)\n\nval added_tokens : t -> is_pair:bool -> int\n(** [added_tokens t ~is_pair] is the number of special tokens [t] adds. Useful\n for calculating the truncation budget. *)\n\n(** {1:fmt Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp] formats a post-processor for inspection. *)\n\n(** {1:serialization Serialization} *)\n\nval of_json : Jsont.json -> (t, string) result\n(** [of_json json] is a post-processor from HuggingFace [tokenizer.json] format.\n Errors if [json] is not an object, has a missing or unknown [\"type\"] field,\n or has invalid parameters. *)\n\nval to_json : t -> Jsont.json\n(** [to_json t] is [t] serialized to HuggingFace [tokenizer.json] format. *)\n"
},
{
"path": "packages/brot/lib/pre_tokenizer.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Types *)\n\ntype behavior =\n [ `Isolated\n | `Removed\n | `Merged_with_previous\n | `Merged_with_next\n | `Contiguous ]\n\ntype prepend_scheme = [ `First | `Never | `Always ]\n\ntype t =\n | Byte_level of {\n add_prefix_space : bool;\n use_regex : bool;\n trim_offsets : bool;\n }\n | Bert\n | Whitespace\n | Whitespace_split\n | Punctuation of { behavior : behavior }\n | Split of { pattern : string; behavior : behavior; invert : bool }\n | Char_delimiter of char\n | Digits of { individual : bool }\n | Metaspace of {\n replacement : char;\n prepend_scheme : prepend_scheme;\n split : bool;\n }\n | Sequence of t list\n | Fixed_length of { length : int }\n | Unicode_scripts\n\n(* Errors *)\n\nlet strf = Printf.sprintf\nlet err_unknown_behavior s = strf \"unknown punctuation behavior '%s'\" s\nlet err_unknown_scheme s = strf \"unknown prepend_scheme '%s'\" s\nlet err_unsupported_type s = strf \"unsupported pre-tokenizer type '%s'\" s\nlet err_expected_char name = strf \"expected single character for '%s'\" name\nlet err_missing_type = \"missing 'type' field\"\nlet err_expected_object = \"expected JSON object\"\nlet err_missing_behavior = \"missing 'behavior' field\"\nlet err_split_missing = \"requires 'pattern' and 'behavior'\"\nlet err_char_delim_missing = \"requires 'delimiter'\"\nlet err_metaspace_missing = \"requires 'replacement' and 'prepend_scheme'\"\nlet err_sequence_missing = \"requires 'pretokenizers' list\"\nlet err_fixed_length = \"requires positive length\"\n\n(* Character classification *)\n\n(* ASCII property table: packed flags for O(1) classification. bit 0:\n whitespace, bit 1: alphabetic, bit 2: numeric, bit 3: punctuation *)\n\nlet ascii_props =\n let t = Array.make 128 0 in\n for i = 9 to 13 do\n t.(i) <- t.(i) lor 1\n done;\n t.(32) <- t.(32) lor 1;\n for i = 65 to 90 do\n t.(i) <- t.(i) lor 2\n done;\n for i = 97 to 122 do\n t.(i) <- t.(i) lor 2\n done;\n for i = 48 to 57 do\n t.(i) <- t.(i) lor 4\n done;\n List.iter\n (fun i -> t.(i) <- t.(i) lor 8)\n [\n 33;\n 34;\n 35;\n 37;\n 38;\n 39;\n 40;\n 41;\n 42;\n 44;\n 45;\n 46;\n 47;\n 58;\n 59;\n 63;\n 64;\n 91;\n 92;\n 93;\n 95;\n 123;\n 125;\n ];\n t\n\nlet[@inline] is_whitespace code =\n if code < 128 then Array.unsafe_get ascii_props code land 1 <> 0\n else Uucp.White.is_white_space (Uchar.of_int code)\n\nlet[@inline] is_alphabetic code =\n if code < 128 then Array.unsafe_get ascii_props code land 2 <> 0\n else Uucp.Alpha.is_alphabetic (Uchar.of_int code)\n\nlet[@inline] is_numeric code =\n if code < 128 then Array.unsafe_get ascii_props code land 4 <> 0\n else\n match Uucp.Gc.general_category (Uchar.of_int code) with\n | `Nd | `Nl | `No -> true\n | _ -> false\n\nlet[@inline] is_punctuation code =\n if code < 128 then Array.unsafe_get ascii_props code land 8 <> 0\n else\n match Uucp.Gc.general_category (Uchar.of_int code) with\n | `Pc | `Pd | `Pe | `Pf | `Pi | `Po | `Ps -> true\n | _ -> false\n\n(* Returns (codepoint lsl 3) lor byte_length — zero allocation. *)\nlet[@inline] utf8_next s i =\n let c = Char.code (String.unsafe_get s i) in\n if c < 0x80 then (c lsl 3) lor 1\n else if c < 0xE0 then\n (((c land 0x1F) lsl 6)\n lor (Char.code (String.unsafe_get s (i + 1)) land 0x3F))\n lsl 3\n lor 2\n else if c < 0xF0 then\n (((c land 0x0F) lsl 12)\n lor ((Char.code (String.unsafe_get s (i + 1)) land 0x3F) lsl 6)\n lor (Char.code (String.unsafe_get s (i + 2)) land 0x3F))\n lsl 3\n lor 3\n else\n (((c land 0x07) lsl 18)\n lor ((Char.code (String.unsafe_get s (i + 1)) land 0x3F) lsl 12)\n lor ((Char.code (String.unsafe_get s (i + 2)) land 0x3F) lsl 6)\n lor (Char.code (String.unsafe_get s (i + 3)) land 0x3F))\n lsl 3\n lor 4\n\n(* Pre-computed byte ↔ unicode mappings for byte-level encode/decode *)\nlet byte_to_unicode, unicode_to_byte =\n let is_direct = Array.make 256 false in\n for i = 33 to 126 do\n is_direct.(i) <- true\n done;\n for i = 161 to 172 do\n is_direct.(i) <- true\n done;\n for i = 174 to 255 do\n is_direct.(i) <- true\n done;\n let byte_to_unicode = Array.make 256 0 in\n let next_code = ref 0 in\n let max_code = ref 0 in\n for b = 0 to 255 do\n let code =\n if is_direct.(b) then b\n else\n let code = 256 + !next_code in\n incr next_code;\n code\n in\n byte_to_unicode.(b) <- code;\n if code > !max_code then max_code := code\n done;\n let unicode_to_byte = Array.make (!max_code + 1) (-1) in\n for b = 0 to 255 do\n let code = byte_to_unicode.(b) in\n if code < Array.length unicode_to_byte then unicode_to_byte.(code) <- b\n done;\n (byte_to_unicode, unicode_to_byte)\n\nlet byte_level_encode text =\n let len = String.length text in\n (* Worst case: every byte remaps to a 2-byte UTF-8 sequence *)\n let result = Bytes.create (len * 2) in\n let j = ref 0 in\n for i = 0 to len - 1 do\n let u =\n Array.unsafe_get byte_to_unicode (Char.code (String.unsafe_get text i))\n in\n if u < 128 then begin\n Bytes.unsafe_set result !j (Char.unsafe_chr u);\n incr j\n end\n else begin\n Bytes.unsafe_set result !j (Char.unsafe_chr (0xC0 lor (u lsr 6)));\n Bytes.unsafe_set result (!j + 1)\n (Char.unsafe_chr (0x80 lor (u land 0x3F)));\n j := !j + 2\n end\n done;\n Bytes.sub_string result 0 !j\n\nlet byte_level_encode_range text ~start ~len =\n let result = Bytes.create (len * 2) in\n let j = ref 0 in\n for i = start to start + len - 1 do\n let u =\n Array.unsafe_get byte_to_unicode (Char.code (String.unsafe_get text i))\n in\n if u < 128 then begin\n Bytes.unsafe_set result !j (Char.unsafe_chr u);\n incr j\n end\n else begin\n Bytes.unsafe_set result !j (Char.unsafe_chr (0xC0 lor (u lsr 6)));\n Bytes.unsafe_set result (!j + 1)\n (Char.unsafe_chr (0x80 lor (u land 0x3F)));\n j := !j + 2\n end\n done;\n Bytes.sub_string result 0 !j\n\nlet byte_level_decode text =\n let len = String.length text in\n let result = Buffer.create len in\n let i = ref 0 in\n while !i < len do\n let b0 = Char.code (String.unsafe_get text !i) in\n if b0 < 128 then begin\n (* ASCII: direct lookup, no utf8_next needed *)\n let byte = Array.unsafe_get unicode_to_byte b0 in\n Buffer.add_char result\n (if byte >= 0 then Char.chr byte else Char.unsafe_chr b0);\n incr i\n end\n else begin\n let p = utf8_next text !i in\n let code = p lsr 3 and clen = p land 7 in\n let byte =\n if code < Array.length unicode_to_byte then unicode_to_byte.(code)\n else -1\n in\n if byte >= 0 then Buffer.add_char result (Char.chr byte)\n else\n for j = !i to !i + clen - 1 do\n Buffer.add_char result (String.unsafe_get text j)\n done;\n i := !i + clen\n end\n done;\n Buffer.contents result\n\nlet[@inline] is_other code =\n (not (is_whitespace code))\n && (not (is_alphabetic code))\n && not (is_numeric code)\n\nlet split_gpt2_pattern text =\n let len = String.length text in\n if len = 0 then []\n else\n let spans = ref [] in\n let pos = ref 0 in\n\n (* Try: optional leading space + run of chars matching a class.\n [ascii_mask]: bitmask into ascii_props for the ASCII fast path. [invert]:\n when true, match chars where (props land mask) = 0. [classify]: predicate\n for non-ASCII codepoints (slow path only). *)\n let try_space_run ~ascii_mask ~invert ~classify () =\n let start = !pos in\n let b0 = Char.code (String.unsafe_get text !pos) in\n let has_space =\n if b0 < 128 then Array.unsafe_get ascii_props b0 land 1 <> 0\n else is_whitespace b0\n in\n let run_start = if has_space then start + 1 else start in\n if run_start < len then\n let b = Char.code (String.unsafe_get text run_start) in\n let ok, clen =\n if b < 128 then\n let v = Array.unsafe_get ascii_props b land ascii_mask in\n ((if invert then v = 0 else v <> 0), 1)\n else\n let p = utf8_next text run_start in\n let code = p lsr 3 and cl = p land 7 in\n (classify code, cl)\n in\n if ok then (\n let j = ref (run_start + clen) in\n let continue = ref true in\n while !j < len && !continue do\n let b = Char.code (String.unsafe_get text !j) in\n if b < 128 then\n let v = Array.unsafe_get ascii_props b land ascii_mask in\n if if invert then v = 0 else v <> 0 then j := !j + 1\n else continue := false\n else\n let p = utf8_next text !j in\n if classify (p lsr 3) then j := !j + (p land 7)\n else continue := false\n done;\n spans := (start, !j - start) :: !spans;\n pos := !j;\n true)\n else false\n else false\n in\n\n let[@inline] next_is_alnum next_pos =\n if next_pos >= len then false\n else\n let nb = Char.code (String.unsafe_get text next_pos) in\n if nb < 128 then Array.unsafe_get ascii_props nb land 6 <> 0\n else\n let nc = utf8_next text next_pos lsr 3 in\n is_alphabetic nc || is_numeric nc\n in\n\n let rec loop () =\n if !pos >= len then ()\n else begin\n (* 1. Contractions: 's 't 'm 'd 're 've 'll *)\n let matched_contraction =\n text.[!pos] = '\\''\n &&\n let remaining = len - !pos in\n remaining >= 2\n &&\n let c1 = String.unsafe_get text (!pos + 1) in\n if c1 = 's' || c1 = 't' || c1 = 'm' || c1 = 'd' then (\n spans := (!pos, 2) :: !spans;\n pos := !pos + 2;\n true)\n else\n remaining >= 3\n &&\n let c2 = String.unsafe_get text (!pos + 2) in\n if\n (c1 = 'r' && c2 = 'e')\n || (c1 = 'v' && c2 = 'e')\n || (c1 = 'l' && c2 = 'l')\n then (\n spans := (!pos, 3) :: !spans;\n pos := !pos + 3;\n true)\n else false\n in\n if matched_contraction then ()\n else if\n try_space_run ~ascii_mask:2 ~invert:false ~classify:is_alphabetic ()\n then ()\n else if\n try_space_run ~ascii_mask:4 ~invert:false ~classify:is_numeric ()\n then ()\n else if try_space_run ~ascii_mask:7 ~invert:true ~classify:is_other ()\n then ()\n (* 5 & 6. Whitespace run *)\n else begin\n let b0 = Char.code (String.unsafe_get text !pos) in\n let is_ws, clen =\n if b0 < 128 then (Array.unsafe_get ascii_props b0 land 1 <> 0, 1)\n else\n let p = utf8_next text !pos in\n let code = p lsr 3 and cl = p land 7 in\n (is_whitespace code, cl)\n in\n if is_ws then begin\n let j = ref (!pos + clen) in\n let continue = ref true in\n while !j < len && !continue do\n let b = Char.code (String.unsafe_get text !j) in\n if b < 128 then\n if Array.unsafe_get ascii_props b land 1 <> 0 then\n if next_is_alnum (!j + 1) && b = 0x20 then continue := false\n else j := !j + 1\n else continue := false\n else\n let p = utf8_next text !j in\n let code = p lsr 3 and cl = p land 7 in\n if is_whitespace code then\n if next_is_alnum (!j + cl) && code = 0x20 then\n continue := false\n else j := !j + cl\n else continue := false\n done;\n spans := (!pos, !j - !pos) :: !spans;\n pos := !j\n end\n else begin\n (* Fallback: single character *)\n spans := (!pos, clen) :: !spans;\n pos := !pos + clen\n end\n end;\n loop ()\n end\n in\n loop ();\n List.rev !spans\n\n(* Pre-tokenize implementations *)\n\nlet pre_tokenize_whitespace_split text =\n let pieces = ref [] in\n let start = ref (-1) in\n let i = ref 0 in\n let len = String.length text in\n let flush () =\n if !start >= 0 then begin\n pieces := (String.sub text !start (!i - !start), (!start, !i)) :: !pieces;\n start := -1\n end\n in\n while !i < len do\n let b = Char.code (String.unsafe_get text !i) in\n if b < 128 then\n if Array.unsafe_get ascii_props b land 1 <> 0 then (\n flush ();\n i := !i + 1)\n else (\n if !start < 0 then start := !i;\n i := !i + 1)\n else\n let p = utf8_next text !i in\n let code = p lsr 3 and l = p land 7 in\n if is_whitespace code then (\n flush ();\n i := !i + l)\n else (\n if !start < 0 then start := !i;\n i := !i + l)\n done;\n flush ();\n List.rev !pieces\n\nlet pre_tokenize_whitespace text =\n let pieces = ref [] in\n let start = ref (-1) in\n let i = ref 0 in\n let len = String.length text in\n let in_word = ref false in\n let in_punct = ref false in\n let flush () =\n if !start >= 0 then begin\n pieces := (String.sub text !start (!i - !start), (!start, !i)) :: !pieces;\n start := -1\n end\n in\n while !i < len do\n let b = Char.code (String.unsafe_get text !i) in\n if b < 128 then\n let p = Array.unsafe_get ascii_props b in\n if p land 6 <> 0 || b = 95 then (\n if !in_punct then flush ();\n if !start < 0 then start := !i;\n in_word := true;\n in_punct := false;\n i := !i + 1)\n else if p land 1 <> 0 then (\n flush ();\n in_word := false;\n in_punct := false;\n i := !i + 1)\n else (\n if !in_word then flush ();\n if !start < 0 then start := !i;\n in_word := false;\n in_punct := true;\n i := !i + 1)\n else\n let p = utf8_next text !i in\n let code = p lsr 3 and l = p land 7 in\n if is_alphabetic code || is_numeric code then (\n if !in_punct then flush ();\n if !start < 0 then start := !i;\n in_word := true;\n in_punct := false;\n i := !i + l)\n else if is_whitespace code then (\n flush ();\n in_word := false;\n in_punct := false;\n i := !i + l)\n else (\n if !in_word then flush ();\n if !start < 0 then start := !i;\n in_word := false;\n in_punct := true;\n i := !i + l)\n done;\n flush ();\n List.rev !pieces\n\nlet pre_tokenize_byte_level ~add_prefix_space ~use_regex ~trim_offsets:_ text =\n let orig_len = String.length text in\n let text, prefix_added =\n if\n add_prefix_space && orig_len > 0\n && not (is_whitespace (Char.code text.[0]))\n then (\" \" ^ text, true)\n else (text, false)\n in\n if use_regex then\n let spans = split_gpt2_pattern text in\n List.map\n (fun (start, plen) ->\n let o_start =\n if prefix_added then if start = 0 then 0 else start - 1 else start\n in\n let o_end =\n min orig_len (if prefix_added then start + plen - 1 else start + plen)\n in\n (byte_level_encode_range text ~start ~len:plen, (max 0 o_start, o_end)))\n spans\n else [ (byte_level_encode text, (0, orig_len)) ]\n\nlet pre_tokenize_bert text =\n let pieces = ref [] in\n let start = ref (-1) in\n let i = ref 0 in\n let len = String.length text in\n let flush () =\n if !start >= 0 then begin\n pieces := (String.sub text !start (!i - !start), (!start, !i)) :: !pieces;\n start := -1\n end\n in\n while !i < len do\n let b = Char.code (String.unsafe_get text !i) in\n if b < 128 then\n let p = Array.unsafe_get ascii_props b in\n if p land 1 <> 0 then (\n flush ();\n i := !i + 1)\n else if p land 8 <> 0 then (\n flush ();\n pieces := (String.sub text !i 1, (!i, !i + 1)) :: !pieces;\n i := !i + 1)\n else (\n if !start < 0 then start := !i;\n i := !i + 1)\n else\n let p = utf8_next text !i in\n let code = p lsr 3 and l = p land 7 in\n if is_whitespace code then (\n flush ();\n i := !i + l)\n else if is_punctuation code then (\n flush ();\n pieces := (String.sub text !i l, (!i, !i + l)) :: !pieces;\n i := !i + l)\n else (\n if !start < 0 then start := !i;\n i := !i + l)\n done;\n flush ();\n List.rev !pieces\n\nlet pre_tokenize_punctuation ~behavior text =\n let pieces = ref [] in\n let start = ref (-1) in\n let i = ref 0 in\n let len = String.length text in\n let last_was_punc = ref false in\n let flush () =\n if !start >= 0 then begin\n pieces := (String.sub text !start (!i - !start), (!start, !i)) :: !pieces;\n start := -1\n end\n in\n let handle_char is_p l =\n if is_p then (\n (match behavior with\n | `Isolated ->\n flush ();\n pieces := (String.sub text !i l, (!i, !i + l)) :: !pieces\n | `Removed -> flush ()\n | `Merged_with_previous -> if !start < 0 then start := !i\n | `Merged_with_next ->\n flush ();\n start := !i\n | `Contiguous ->\n if not (!start >= 0 && !last_was_punc) then begin\n flush ();\n start := !i\n end);\n last_was_punc := true;\n i := !i + l)\n else (\n if behavior = `Contiguous && !start >= 0 && !last_was_punc then flush ();\n if !start < 0 then start := !i;\n i := !i + l;\n last_was_punc := false)\n in\n while !i < len do\n let b = Char.code (String.unsafe_get text !i) in\n if b < 128 then handle_char (Array.unsafe_get ascii_props b land 8 <> 0) 1\n else\n let p = utf8_next text !i in\n let code = p lsr 3 and l = p land 7 in\n handle_char (is_punctuation code) l\n done;\n flush ();\n List.rev !pieces\n\nlet pre_tokenize_split ~pattern ~behavior ~invert text =\n let plen = String.length pattern in\n if plen = 0 then [ (text, (0, String.length text)) ]\n else\n let pieces = ref [] in\n let current = Buffer.create 16 in\n let current_start = ref 0 in\n let i = ref 0 in\n let flush_current () =\n if Buffer.length current > 0 then (\n pieces :=\n ( Buffer.contents current,\n (!current_start, !current_start + Buffer.length current) )\n :: !pieces;\n Buffer.clear current)\n in\n while !i < String.length text do\n let is_match =\n !i + plen <= String.length text && String.sub text !i plen = pattern\n in\n let is_delim = if invert then not is_match else is_match in\n let delim_len = if is_delim then if invert then 1 else plen else 1 in\n if is_delim then (\n (match behavior with\n | `Removed -> flush_current ()\n | `Isolated ->\n flush_current ();\n let delim_str = String.sub text !i delim_len in\n pieces := (delim_str, (!i, !i + delim_len)) :: !pieces\n | `Merged_with_previous ->\n Buffer.add_string current (String.sub text !i delim_len);\n flush_current ()\n | `Merged_with_next ->\n flush_current ();\n current_start := !i;\n Buffer.add_string current (String.sub text !i delim_len)\n | `Contiguous ->\n if Buffer.length current > 0 && is_delim then\n Buffer.add_string current (String.sub text !i delim_len)\n else (\n flush_current ();\n Buffer.add_string current (String.sub text !i delim_len)));\n i := !i + delim_len)\n else (\n if Buffer.length current = 0 then current_start := !i;\n Buffer.add_string current (String.sub text !i 1);\n i := !i + 1)\n done;\n flush_current ();\n List.rev !pieces\n\nlet pre_tokenize_digits ~individual text =\n let pieces = ref [] in\n let start = ref (-1) in\n let i = ref 0 in\n let len = String.length text in\n let in_digits = ref false in\n let flush () =\n if !start >= 0 then begin\n pieces := (String.sub text !start (!i - !start), (!start, !i)) :: !pieces;\n start := -1\n end\n in\n let handle_char is_d l =\n if individual && is_d then (\n flush ();\n pieces := (String.sub text !i l, (!i, !i + l)) :: !pieces;\n i := !i + l)\n else (\n if is_d <> !in_digits then (\n flush ();\n in_digits := is_d);\n if !start < 0 then start := !i;\n i := !i + l)\n in\n while !i < len do\n let b = Char.code (String.unsafe_get text !i) in\n if b < 128 then handle_char (Array.unsafe_get ascii_props b land 4 <> 0) 1\n else\n let p = utf8_next text !i in\n let code = p lsr 3 and l = p land 7 in\n handle_char (is_numeric code) l\n done;\n flush ();\n List.rev !pieces\n\nlet pre_tokenize_metaspace ~replacement ~prepend_scheme ~split text =\n let repl = String.make 1 replacement in\n let text =\n match prepend_scheme with\n | (`Always | `First) when String.length text > 0 && text.[0] <> ' ' ->\n \" \" ^ text\n | _ -> text\n in\n let len = String.length text in\n let buf = Buffer.create len in\n let i = ref 0 in\n while !i < len do\n if text.[!i] = ' ' then (\n Buffer.add_string buf repl;\n incr i)\n else\n let l = utf8_next text !i land 7 in\n Buffer.add_substring buf text !i l;\n i := !i + l\n done;\n let transformed = Buffer.contents buf in\n if split then (\n let tlen = String.length transformed in\n let rlen = String.length repl in\n let splits = ref [] in\n let start = ref 0 in\n let pos = ref 0 in\n while !pos < tlen do\n if !pos + rlen <= tlen && String.sub transformed !pos rlen = repl then (\n if !pos > !start then\n splits :=\n (String.sub transformed !start (!pos - !start), (!start, !pos))\n :: !splits;\n start := !pos;\n pos := !pos + rlen)\n else incr pos\n done;\n if !pos > !start then\n splits :=\n (String.sub transformed !start (!pos - !start), (!start, !pos))\n :: !splits;\n List.rev !splits)\n else [ (transformed, (0, len)) ]\n\nlet pre_tokenize_fixed_length ~length text =\n if length <= 0 || String.length text = 0 then []\n else\n let pieces = ref [] in\n let len = String.length text in\n let i = ref 0 in\n while !i < len do\n let start = !i in\n let count = ref 0 in\n while !i < len && !count < length do\n let l = utf8_next text !i land 7 in\n i := !i + l;\n incr count\n done;\n pieces := (String.sub text start (!i - start), (start, !i)) :: !pieces\n done;\n List.rev !pieces\n\ntype script = [ `Any | Uucp.Script.t ]\n\nlet fixed_script code : script =\n if code = 0x30FC then (`Hani :> script)\n else if is_whitespace code then `Any\n else\n match Uucp.Script.script (Uchar.of_int code) with\n | `Hira | `Kana -> (`Hani :> script)\n | s -> (s :> script)\n\nlet pre_tokenize_unicode_scripts text =\n let pieces = ref [] in\n let start = ref (-1) in\n let len = String.length text in\n let i = ref 0 in\n let last_script = ref None in\n let flush () =\n if !start >= 0 then begin\n pieces := (String.sub text !start (!i - !start), (!start, !i)) :: !pieces;\n start := -1\n end\n in\n let emit (script : script) l =\n if\n script <> `Any && !last_script <> Some `Any && !last_script <> Some script\n then flush ();\n if !start < 0 then start := !i;\n i := !i + l;\n if script <> `Any then last_script := Some script\n in\n while !i < len do\n let b = Char.code (String.unsafe_get text !i) in\n if b < 128 then\n let p = Array.unsafe_get ascii_props b in\n let script : script =\n if p land 1 <> 0 then `Any else if p land 2 <> 0 then `Latn else `Zyyy\n in\n emit script 1\n else\n let p = utf8_next text !i in\n let code = p lsr 3 and l = p land 7 in\n emit (fixed_script code) l\n done;\n flush ();\n List.rev !pieces\n\n(* Constructors *)\n\nlet whitespace () = Whitespace\nlet whitespace_split () = Whitespace_split\nlet bert () = Bert\n\nlet byte_level ?(add_prefix_space = true) ?(use_regex = true)\n ?(trim_offsets = true) () =\n Byte_level { add_prefix_space; use_regex; trim_offsets }\n\nlet punctuation ?(behavior = `Isolated) () = Punctuation { behavior }\n\nlet split ~pattern ?(behavior = `Removed) ?(invert = false) () =\n Split { pattern; behavior; invert }\n\nlet char_delimiter c = Char_delimiter c\n\nlet digits ?(individual_digits = false) () =\n Digits { individual = individual_digits }\n\nlet metaspace ?(replacement = '_') ?(prepend_scheme = `Always) ?(split = true)\n () =\n Metaspace { replacement; prepend_scheme; split }\n\nlet unicode_scripts () = Unicode_scripts\nlet fixed_length n = Fixed_length { length = n }\nlet sequence ts = Sequence ts\n\n(* Dispatch *)\n\nlet rec pre_tokenize t text =\n match t with\n | Whitespace -> pre_tokenize_whitespace text\n | Whitespace_split -> pre_tokenize_whitespace_split text\n | Bert -> pre_tokenize_bert text\n | Byte_level { add_prefix_space; use_regex; trim_offsets } ->\n pre_tokenize_byte_level ~add_prefix_space ~use_regex ~trim_offsets text\n | Punctuation { behavior } -> pre_tokenize_punctuation ~behavior text\n | Split { pattern; behavior; invert } ->\n pre_tokenize_split ~pattern ~behavior ~invert text\n | Char_delimiter c ->\n pre_tokenize_split ~pattern:(String.make 1 c) ~behavior:`Removed\n ~invert:false text\n | Digits { individual } -> pre_tokenize_digits ~individual text\n | Metaspace { replacement; prepend_scheme; split } ->\n pre_tokenize_metaspace ~replacement ~prepend_scheme ~split text\n | Unicode_scripts -> pre_tokenize_unicode_scripts text\n | Fixed_length { length } -> pre_tokenize_fixed_length ~length text\n | Sequence ts -> pre_tokenize_sequence ts text\n\nand pre_tokenize_sequence ts text =\n let initial = [ (text, (0, String.length text)) ] in\n List.fold_left\n (fun pieces t ->\n List.concat_map\n (fun (s, (o_start, _)) ->\n let sub_pieces = pre_tokenize t s in\n List.map\n (fun (p, (p_start, p_end)) ->\n (p, (o_start + p_start, o_start + p_end)))\n sub_pieces)\n pieces)\n initial ts\n\n(* Serialization *)\n\nlet json_obj pairs =\n Jsont.Json.object' (List.map (fun (k, v) -> (Jsont.Json.name k, v)) pairs)\n\nlet behavior_to_string = function\n | `Isolated -> \"Isolated\"\n | `Removed -> \"Removed\"\n | `Merged_with_previous -> \"MergedWithPrevious\"\n | `Merged_with_next -> \"MergedWithNext\"\n | `Contiguous -> \"Contiguous\"\n\nlet behavior_of_string = function\n | \"Isolated\" -> Ok `Isolated\n | \"Removed\" -> Ok `Removed\n | \"MergedWithPrevious\" -> Ok `Merged_with_previous\n | \"MergedWithNext\" -> Ok `Merged_with_next\n | \"Contiguous\" -> Ok `Contiguous\n | other -> Error (err_unknown_behavior other)\n\nlet scheme_to_string = function\n | `First -> \"First\"\n | `Never -> \"Never\"\n | `Always -> \"Always\"\n\nlet scheme_of_string = function\n | \"First\" -> Ok `First\n | \"Never\" -> Ok `Never\n | \"Always\" -> Ok `Always\n | other -> Error (err_unknown_scheme other)\n\n(* Formatting *)\n\nlet rec pp ppf = function\n | Byte_level { add_prefix_space; use_regex; trim_offsets } ->\n Format.fprintf ppf\n \"@[<1>ByteLevel(add_prefix_space=%b,@ use_regex=%b,@ trim_offsets=%b)@]\"\n add_prefix_space use_regex trim_offsets\n | Bert -> Format.pp_print_string ppf \"Bert\"\n | Whitespace -> Format.pp_print_string ppf \"Whitespace\"\n | Whitespace_split -> Format.pp_print_string ppf \"WhitespaceSplit\"\n | Punctuation { behavior } ->\n Format.fprintf ppf \"@[<1>Punctuation(%s)@]\" (behavior_to_string behavior)\n | Split { pattern; behavior; invert } ->\n Format.fprintf ppf \"@[<1>Split(%S,@ %s,@ invert=%b)@]\" pattern\n (behavior_to_string behavior)\n invert\n | Char_delimiter c -> Format.fprintf ppf \"CharDelimiter(%C)\" c\n | Digits { individual } ->\n Format.fprintf ppf \"Digits(individual=%b)\" individual\n | Metaspace { replacement; prepend_scheme; split } ->\n Format.fprintf ppf \"@[<1>Metaspace(%C,@ %s,@ split=%b)@]\" replacement\n (scheme_to_string prepend_scheme)\n split\n | Sequence ts ->\n Format.fprintf ppf \"@[<1>Sequence[%a]@]\"\n (Format.pp_print_list\n ~pp_sep:(fun ppf () -> Format.fprintf ppf \",@ \")\n pp)\n ts\n | Fixed_length { length } -> Format.fprintf ppf \"FixedLength(%d)\" length\n | Unicode_scripts -> Format.pp_print_string ppf \"UnicodeScripts\"\n\nlet rec to_json = function\n | Byte_level { add_prefix_space; use_regex; trim_offsets } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"ByteLevel\");\n (\"add_prefix_space\", Jsont.Json.bool add_prefix_space);\n (\"use_regex\", Jsont.Json.bool use_regex);\n (\"trim_offsets\", Jsont.Json.bool trim_offsets);\n ]\n | Bert -> json_obj [ (\"type\", Jsont.Json.string \"BertPreTokenizer\") ]\n | Whitespace -> json_obj [ (\"type\", Jsont.Json.string \"Whitespace\") ]\n | Whitespace_split ->\n json_obj [ (\"type\", Jsont.Json.string \"WhitespaceSplit\") ]\n | Punctuation { behavior } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Punctuation\");\n (\"behavior\", Jsont.Json.string (behavior_to_string behavior));\n ]\n | Split { pattern; behavior; invert } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Split\");\n (\"pattern\", Jsont.Json.string pattern);\n (\"behavior\", Jsont.Json.string (behavior_to_string behavior));\n (\"invert\", Jsont.Json.bool invert);\n ]\n | Char_delimiter delimiter ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"CharDelimiterSplit\");\n (\"delimiter\", Jsont.Json.string (String.make 1 delimiter));\n ]\n | Digits { individual } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Digits\");\n (\"individual_digits\", Jsont.Json.bool individual);\n ]\n | Metaspace { replacement; prepend_scheme; split } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Metaspace\");\n (\"replacement\", Jsont.Json.string (String.make 1 replacement));\n (\"prepend_scheme\", Jsont.Json.string (scheme_to_string prepend_scheme));\n (\"split\", Jsont.Json.bool split);\n ]\n | Sequence ts ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"Sequence\");\n (\"pretokenizers\", Jsont.Json.list (List.map to_json ts));\n ]\n | Fixed_length { length } ->\n json_obj\n [\n (\"type\", Jsont.Json.string \"FixedLength\");\n (\"length\", Jsont.Json.int length);\n ]\n | Unicode_scripts -> json_obj [ (\"type\", Jsont.Json.string \"UnicodeScripts\") ]\n\nlet find_field name fields = Option.map snd (Jsont.Json.find_mem name fields)\n\nlet bool_field name default fields =\n match find_field name fields with\n | Some (Jsont.Bool (b, _)) -> b\n | Some (Jsont.Number (f, _)) -> int_of_float f <> 0\n | Some (Jsont.String (s, _)) -> (\n match String.lowercase_ascii s with\n | \"true\" | \"1\" -> true\n | \"false\" | \"0\" -> false\n | _ -> default)\n | _ -> default\n\nlet int_field name default fields =\n match find_field name fields with\n | Some (Jsont.Number (f, _)) -> int_of_float f\n | Some (Jsont.String (s, _)) -> (\n match int_of_string_opt s with Some v -> v | None -> default)\n | _ -> default\n\nlet char_of_field name = function\n | Jsont.String (s, _) when String.length s = 1 -> Ok s.[0]\n | _ -> Error (err_expected_char name)\n\nlet rec of_json = function\n | Jsont.Object (fields, _) -> (\n match find_field \"type\" fields with\n | Some (Jsont.String (\"ByteLevel\", _)) ->\n let add_prefix_space = bool_field \"add_prefix_space\" true fields in\n let use_regex = bool_field \"use_regex\" true fields in\n let trim_offsets = bool_field \"trim_offsets\" true fields in\n Ok (Byte_level { add_prefix_space; use_regex; trim_offsets })\n | Some (Jsont.String (\"BertPreTokenizer\", _)) -> Ok Bert\n | Some (Jsont.String (\"Whitespace\", _)) -> Ok Whitespace\n | Some (Jsont.String (\"WhitespaceSplit\", _)) -> Ok Whitespace_split\n | Some (Jsont.String (\"Punctuation\", _)) -> (\n match find_field \"behavior\" fields with\n | Some (Jsont.String (s, _)) ->\n Result.map\n (fun b -> Punctuation { behavior = b })\n (behavior_of_string s)\n | _ -> Error err_missing_behavior)\n | Some (Jsont.String (\"Split\", _)) -> (\n match (find_field \"pattern\" fields, find_field \"behavior\" fields) with\n | ( Some (Jsont.String (pattern, _)),\n Some (Jsont.String (behavior_str, _)) ) ->\n Result.map\n (fun behavior ->\n let invert = bool_field \"invert\" false fields in\n Split { pattern; behavior; invert })\n (behavior_of_string behavior_str)\n | _ -> Error err_split_missing)\n | Some (Jsont.String (\"CharDelimiterSplit\", _)) -> (\n match find_field \"delimiter\" fields with\n | Some v ->\n Result.map\n (fun c -> Char_delimiter c)\n (char_of_field \"delimiter\" v)\n | None -> Error err_char_delim_missing)\n | Some (Jsont.String (\"Digits\", _)) ->\n let individual = bool_field \"individual_digits\" false fields in\n Ok (Digits { individual })\n | Some (Jsont.String (\"Metaspace\", _)) -> (\n match\n (find_field \"replacement\" fields, find_field \"prepend_scheme\" fields)\n with\n | Some (Jsont.String (repl, _)), Some (Jsont.String (scheme, _))\n when String.length repl = 1 ->\n Result.map\n (fun prepend_scheme ->\n let split = bool_field \"split\" true fields in\n Metaspace { replacement = repl.[0]; prepend_scheme; split })\n (scheme_of_string scheme)\n | _ -> Error err_metaspace_missing)\n | Some (Jsont.String (\"Sequence\", _)) -> (\n match find_field \"pretokenizers\" fields with\n | Some (Jsont.Array (elements, _)) ->\n let rec build acc = function\n | [] -> Ok (Sequence (List.rev acc))\n | item :: rest -> (\n match of_json item with\n | Ok t -> build (t :: acc) rest\n | Error _ as e -> e)\n in\n build [] elements\n | _ -> Error err_sequence_missing)\n | Some (Jsont.String (\"FixedLength\", _)) ->\n let length = int_field \"length\" 0 fields in\n if length <= 0 then Error err_fixed_length\n else Ok (Fixed_length { length })\n | Some (Jsont.String (\"UnicodeScripts\", _)) -> Ok Unicode_scripts\n | Some (Jsont.String (other, _)) -> Error (err_unsupported_type other)\n | _ -> Error err_missing_type)\n | _ -> Error err_expected_object\n"
},
{
"path": "packages/brot/lib/pre_tokenizer.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Pre-tokenization.\n\n Pre-tokenizers split raw text into pieces before vocabulary-based\n tokenization (BPE, WordPiece, etc.) is applied. Each piece carries byte\n offsets into the original text.\n\n See {!Brot} for the full tokenization pipeline. *)\n\ntype t\n(** The type for pre-tokenizers. *)\n\n(** {1:constructors Constructors} *)\n\nval whitespace : unit -> t\n(** [whitespace ()] splits on whitespace using pattern [\\w+|[^\\w\\s]+].\n\n Groups word characters (letters, digits, underscore) together and groups\n non-word, non-space characters together. Whitespace is used as delimiter but\n not included in output. *)\n\nval whitespace_split : unit -> t\n(** [whitespace_split ()] splits on any whitespace characters.\n\n Removes whitespace from output. Simplest and fastest pre-tokenizer. *)\n\nval bert : unit -> t\n(** [bert ()] applies BERT-style pre-tokenization.\n\n Splits on whitespace, isolates punctuation, and separates CJK characters\n individually. *)\n\nval byte_level :\n ?add_prefix_space:bool -> ?use_regex:bool -> ?trim_offsets:bool -> unit -> t\n(** [byte_level ()] is a byte-level pre-tokenizer. Used by GPT-2, GPT-3,\n RoBERTa.\n\n Converts text to byte representation and applies GPT-2's regex pattern for\n splitting.\n\n - [add_prefix_space]: add space at beginning if text does not start with\n whitespace. Default: [true].\n - [use_regex]: use GPT-2's regex pattern for splitting. Default: [true].\n - [trim_offsets]: adjust offsets for byte-level encoding. Default: [true].\n*)\n\ntype behavior =\n [ `Isolated (** Keep delimiter as separate piece *)\n | `Removed (** Remove delimiter *)\n | `Merged_with_previous (** Merge delimiter with previous piece *)\n | `Merged_with_next (** Merge delimiter with next piece *)\n | `Contiguous (** Group consecutive delimiters together *) ]\n(** Delimiter handling behavior for splitting operations. *)\n\nval punctuation : ?behavior:behavior -> unit -> t\n(** [punctuation ()] separates punctuation from alphanumeric content.\n\n [behavior] defaults to [`Isolated]. *)\n\nval split : pattern:string -> ?behavior:behavior -> ?invert:bool -> unit -> t\n(** [split ~pattern ()] splits on a literal string [pattern].\n\n [behavior] defaults to [`Removed]. When [invert] is [true], splits on\n everything {e except} the pattern; defaults to [false]. *)\n\nval char_delimiter : char -> t\n(** [char_delimiter c] splits on character [c], removing it from output.\n\n Equivalent to [split ~pattern:(String.make 1 c) ~behavior:`Removed ()]. *)\n\nval digits : ?individual_digits:bool -> unit -> t\n(** [digits ()] splits on digit boundaries.\n\n When [individual_digits] is [true], each digit is a separate piece; when\n [false] (default), consecutive digits are grouped. *)\n\ntype prepend_scheme =\n [ `First (** Only prepend to first piece *)\n | `Never (** Never prepend *)\n | `Always (** Always prepend if not starting with space *) ]\n(** Controls when metaspace prepends the replacement character. *)\n\nval metaspace :\n ?replacement:char ->\n ?prepend_scheme:prepend_scheme ->\n ?split:bool ->\n unit ->\n t\n(** [metaspace ()] replaces whitespace with a visible marker. Used by\n SentencePiece models.\n\n - [replacement]: character to replace spaces with. Default: ['_'].\n - [prepend_scheme]: when to prepend the replacement character. Default:\n [`Always].\n - [split]: whether to split on the replacement character. Default: [true].\n*)\n\nval unicode_scripts : unit -> t\n(** [unicode_scripts ()] splits on Unicode script boundaries.\n\n Separates text when the writing system changes (e.g., Latin to Cyrillic,\n Latin to Han). *)\n\nval fixed_length : int -> t\n(** [fixed_length n] splits into fixed-length character chunks.\n\n The last chunk may be shorter than [n]. *)\n\nval sequence : t list -> t\n(** [sequence ts] chains multiple pre-tokenizers left-to-right.\n\n Each pre-tokenizer processes the pieces from the previous one. Offsets are\n composed correctly through the chain. *)\n\n(** {1 Operations} *)\n\nval pre_tokenize : t -> string -> (string * (int * int)) list\n(** [pre_tokenize t text] splits [text] into pieces with character offsets.\n\n Returns a list of [(piece, (start, end_))] where [start] and [end_] are byte\n positions in the original [text]. Offsets are non-overlapping and in\n ascending order. *)\n\n(** {1 Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp ppf t] formats [t] for inspection. *)\n\n(** {1:byte_level_decode Byte-level decoding} *)\n\nval byte_level_decode : string -> string\n(** [byte_level_decode s] reverses byte-level encoding by converting the special\n Unicode codepoints back to original byte values. *)\n\n(** {1 Serialization} *)\n\nval to_json : t -> Jsont.json\n(** [to_json t] serializes [t] to HuggingFace JSON format. *)\n\nval of_json : Jsont.json -> (t, string) result\n(** [of_json json] is a pre-tokenizer from HuggingFace JSON format. Errors if\n [json] is not an object, has a missing or unknown [\"type\"] field, or has\n invalid parameters. *)\n"
},
{
"path": "packages/brot/lib/unigram.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Compact trie for longest-prefix matching *)\n\ntype trie = {\n trie_ids : int array;\n child_starts : int array;\n edge_bytes : bytes;\n edge_targets : int array;\n}\n\nlet build_trie token_to_ids =\n if Hashtbl.length token_to_ids = 0 then\n {\n trie_ids = [||];\n child_starts = [| 0 |];\n edge_bytes = Bytes.empty;\n edge_targets = [||];\n }\n else\n let cap = ref 256 in\n let ids = ref (Array.make !cap (-1)) in\n let ch = ref (Array.init !cap (fun _ -> Hashtbl.create 0)) in\n let n = ref 1 in\n !ch.(0) <- Hashtbl.create 64;\n let grow () =\n let new_cap = !cap * 2 in\n let new_ids = Array.make new_cap (-1) in\n Array.blit !ids 0 new_ids 0 !n;\n ids := new_ids;\n let new_ch =\n Array.init new_cap (fun i ->\n if i < !n then !ch.(i) else Hashtbl.create 0)\n in\n ch := new_ch;\n cap := new_cap\n in\n Hashtbl.iter\n (fun key id ->\n let cur = ref 0 in\n for i = 0 to String.length key - 1 do\n let byte = Char.code (String.unsafe_get key i) in\n let child =\n match Hashtbl.find_opt !ch.(!cur) byte with\n | Some c -> c\n | None ->\n if !n >= !cap then grow ();\n let c = !n in\n incr n;\n !ch.(c) <- Hashtbl.create 4;\n Hashtbl.add !ch.(!cur) byte c;\n c\n in\n cur := child\n done;\n !ids.(!cur) <- id)\n token_to_ids;\n let node_count = !n in\n let trie_ids = Array.init node_count (fun i -> !ids.(i)) in\n let child_starts = Array.make (node_count + 1) 0 in\n let total = ref 0 in\n for i = 0 to node_count - 1 do\n child_starts.(i) <- !total;\n total := !total + Hashtbl.length !ch.(i)\n done;\n child_starts.(node_count) <- !total;\n let edge_bytes = Bytes.create !total in\n let edge_targets = Array.make !total 0 in\n let pos = ref 0 in\n for i = 0 to node_count - 1 do\n Hashtbl.iter\n (fun byte child ->\n Bytes.unsafe_set edge_bytes !pos (Char.unsafe_chr byte);\n edge_targets.(!pos) <- child;\n incr pos)\n !ch.(i)\n done;\n for i = 0 to node_count - 1 do\n let start = child_starts.(i) in\n let stop = child_starts.(i + 1) in\n for j = start + 1 to stop - 1 do\n let kb = Bytes.unsafe_get edge_bytes j in\n let kt = edge_targets.(j) in\n let k = ref (j - 1) in\n while !k >= start && Bytes.unsafe_get edge_bytes !k > kb do\n Bytes.unsafe_set edge_bytes (!k + 1) (Bytes.unsafe_get edge_bytes !k);\n edge_targets.(!k + 1) <- edge_targets.(!k);\n decr k\n done;\n Bytes.unsafe_set edge_bytes (!k + 1) kb;\n edge_targets.(!k + 1) <- kt\n done\n done;\n { trie_ids; child_starts; edge_bytes; edge_targets }\n\nlet[@inline] trie_step trie node byte =\n let lo = ref (Array.unsafe_get trie.child_starts node) in\n let hi = ref (Array.unsafe_get trie.child_starts (node + 1) - 1) in\n let result = ref (-1) in\n while !lo <= !hi do\n let mid = !lo + ((!hi - !lo) asr 1) in\n let mid_byte = Char.code (Bytes.unsafe_get trie.edge_bytes mid) in\n if mid_byte = byte then (\n result := Array.unsafe_get trie.edge_targets mid;\n lo := !hi + 1)\n else if mid_byte < byte then lo := mid + 1\n else hi := mid - 1\n done;\n !result\n\nlet trie_longest_match trie text ~start =\n if Array.length trie.trie_ids = 0 then None\n else\n let text_len = String.length text in\n let last_id = ref (-1) in\n let last_end = ref start in\n let current = ref 0 in\n let stopped = ref false in\n let j = ref start in\n while !j < text_len && not !stopped do\n let child =\n trie_step trie !current (Char.code (String.unsafe_get text !j))\n in\n if child < 0 then stopped := true\n else (\n current := child;\n incr j;\n let tid = Array.unsafe_get trie.trie_ids child in\n if tid >= 0 then (\n last_id := tid;\n last_end := !j))\n done;\n if !last_id >= 0 then Some (!last_id, !last_end) else None\n\n(* Model type *)\n\ntype t = {\n vocab : (string * float) array;\n token_to_ids : (string, int) Hashtbl.t;\n trie : trie;\n}\n\nlet create vocab_list =\n let vocab = Array.of_list vocab_list in\n let token_to_ids = Hashtbl.create (Array.length vocab) in\n Array.iteri\n (fun idx (token, _) -> Hashtbl.replace token_to_ids token idx)\n vocab;\n let trie = build_trie token_to_ids in\n { vocab; token_to_ids; trie }\n\nlet token_to_id model token = Hashtbl.find_opt model.token_to_ids token\n\nlet id_to_token model id =\n if id >= 0 && id < Array.length model.vocab then\n let token, _ = model.vocab.(id) in\n Some token\n else None\n\nlet get_vocab model = Array.to_list model.vocab\nlet get_vocab_size model = Array.length model.vocab\n\nlet tokenize model text =\n let len = String.length text in\n let rec consume pos acc =\n if pos >= len then List.rev acc\n else if\n text.[pos] = ' '\n || text.[pos] = '\\n'\n || text.[pos] = '\\t'\n || text.[pos] = '\\r'\n then consume (pos + 1) acc\n else\n match trie_longest_match model.trie text ~start:pos with\n | Some (id, end_pos) ->\n let s = String.sub text pos (end_pos - pos) in\n consume end_pos ((id, s, (pos, end_pos)) :: acc)\n | None ->\n let s = String.sub text pos 1 in\n let id = match token_to_id model s with Some id -> id | None -> 0 in\n consume (pos + 1) ((id, s, (pos, pos + 1)) :: acc)\n in\n consume 0 []\n\nlet json_obj pairs =\n Jsont.Json.object' (List.map (fun (k, v) -> (Jsont.Json.name k, v)) pairs)\n\nlet json_to_string j =\n match Jsont_bytesrw.encode_string ~format:Jsont.Minify Jsont.json j with\n | Ok s -> s\n | Error e -> failwith e\n\nlet save model ~folder () =\n let json_vocab =\n Array.to_list model.vocab\n |> List.mapi (fun id (token, prob) ->\n json_obj\n [\n (\"id\", Jsont.Json.int id);\n (\"token\", Jsont.Json.string token);\n (\"prob\", Jsont.Json.number prob);\n ])\n in\n let json =\n json_obj\n [\n (\"type\", Jsont.Json.string \"Unigram\");\n (\"vocab\", Jsont.Json.list json_vocab);\n ]\n in\n let path = Filename.concat folder \"unigram.json\" in\n let oc = open_out path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc (json_to_string json));\n [ \"unigram.json\" ]\n\nlet train ~vocab_size ~show_progress ~special_tokens ~shrinking_factor\n ~unk_token ~max_piece_length ~n_sub_iterations texts existing =\n let _ =\n ( show_progress,\n shrinking_factor,\n unk_token,\n max_piece_length,\n n_sub_iterations,\n existing )\n in\n let counts = Hashtbl.create 10000 in\n List.iter\n (fun line ->\n let words = Re.split (Re.compile (Re.rep1 (Re.set \" \\t\\n\\r\"))) line in\n List.iter\n (fun word ->\n if word <> \"\" then\n Hashtbl.replace counts word\n (1 + Option.value ~default:0 (Hashtbl.find_opt counts word)))\n words)\n texts;\n\n let total =\n Hashtbl.fold (fun _ count acc -> acc + count) counts 0 |> float_of_int\n in\n let sorted =\n Hashtbl.fold (fun token count acc -> (token, count) :: acc) counts []\n |> List.sort (fun (_, c1) (_, c2) -> compare c2 c1)\n in\n\n let take_first n lst =\n let rec aux i = function\n | [] -> []\n | _ when i = 0 -> []\n | x :: xs -> x :: aux (i - 1) xs\n in\n aux n lst\n in\n\n let selected = take_first vocab_size sorted in\n let vocab_with_probs =\n special_tokens\n |> List.map (fun token -> (token, 1.0 /. float_of_int (vocab_size + 1)))\n |> fun specials ->\n specials\n @ List.map\n (fun (token, count) ->\n let prob = if total = 0. then 0. else float_of_int count /. total in\n (token, prob))\n selected\n in\n let model = create vocab_with_probs in\n (model, special_tokens)\n"
},
{
"path": "packages/brot/lib/unigram.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Unigram language model tokenization.\n\n {b Internal module.} Probabilistic subword tokenization using token\n log-probabilities. Used by SentencePiece, AlBERT, T5, and mBART.\n\n Tokenization uses greedy longest-prefix matching via a compact trie with\n sorted edges and binary-search dispatch. At each byte position the longest\n vocabulary match is consumed. Unknown single characters default to ID [0].\n*)\n\ntype t\n(** The type for unigram models. *)\n\n(** {1:creation Creation} *)\n\nval create : (string * float) list -> t\n(** [create vocab] is a unigram model from [(token, log_probability)] pairs. The\n trie is built at creation time. *)\n\n(** {1:tokenization Tokenization} *)\n\nval tokenize : t -> string -> (int * string * (int * int)) list\n(** [tokenize t s] is the tokenization of [s] as [(id, token, (start, stop))]\n triples. Offsets are byte positions in [s]. *)\n\n(** {1:vocabulary Vocabulary} *)\n\nval token_to_id : t -> string -> int option\n(** [token_to_id t tok] is the ID of [tok] in the vocabulary. *)\n\nval id_to_token : t -> int -> string option\n(** [id_to_token t id] is the token string for [id]. *)\n\nval get_vocab : t -> (string * float) list\n(** [get_vocab t] is the vocabulary as [(token, score)] pairs. *)\n\nval get_vocab_size : t -> int\n(** [get_vocab_size t] is the number of tokens in the vocabulary. *)\n\n(** {1:serialization Serialization} *)\n\nval save : t -> folder:string -> unit -> string list\n(** [save t ~folder ()] writes [unigram.json] to [folder]. The file contains\n each token with its ID and log-probability in JSON format. Returns the list\n of created filenames. *)\n\n(** {1:training Training} *)\n\nval train :\n vocab_size:int ->\n show_progress:bool ->\n special_tokens:string list ->\n shrinking_factor:float ->\n unk_token:string option ->\n max_piece_length:int ->\n n_sub_iterations:int ->\n string list ->\n t option ->\n t * string list\n(** [train ~vocab_size ~show_progress ~special_tokens ~shrinking_factor\n ~unk_token ~max_piece_length ~n_sub_iterations texts init] learns a unigram\n model from [texts].\n\n - [vocab_size] is the target vocabulary size.\n - [show_progress] enables progress output on [stderr].\n - [special_tokens] are added to the vocabulary first.\n - [shrinking_factor] controls vocabulary pruning rate.\n - [unk_token] is the unknown token, if any.\n - [max_piece_length] limits the byte length of vocabulary pieces.\n - [n_sub_iterations] is the number of EM sub-iterations.\n - [init], when provided, seeds the vocabulary from an existing model.\n\n Returns [(model, special_tokens)]. *)\n"
},
{
"path": "packages/brot/lib/word_level.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype t = {\n vocab : (string, int) Hashtbl.t;\n vocab_r : (int, string) Hashtbl.t;\n unk_token : string;\n}\n\nlet create ?(vocab = []) ?(unk_token = \"\") () =\n let size = max 1 (List.length vocab) in\n let vocab_tbl = Hashtbl.create size in\n let vocab_r_tbl = Hashtbl.create size in\n List.iter\n (fun (token, id) ->\n Hashtbl.replace vocab_tbl token id;\n Hashtbl.replace vocab_r_tbl id token)\n vocab;\n { vocab = vocab_tbl; vocab_r = vocab_r_tbl; unk_token }\n\nlet add_token vocab vocab_r token id =\n Hashtbl.replace vocab token id;\n Hashtbl.replace vocab_r id token\n\nlet tokenize model text =\n if String.length text = 0 then []\n else\n (* Match HuggingFace tokenizers semantics exactly: 1. Try to find token in\n vocab 2. Fall back to UNK token if available 3. Return empty list if\n neither exists (error case) *)\n match Hashtbl.find_opt model.vocab text with\n | Some id -> [ (id, text, (0, String.length text)) ]\n | None -> (\n match Hashtbl.find_opt model.vocab model.unk_token with\n | Some unk_id -> [ (unk_id, model.unk_token, (0, String.length text)) ]\n | None -> [] (* Token not found and no UNK token - return empty *))\n\nlet tokenize_ids model text =\n if String.length text = 0 then [||]\n else\n match Hashtbl.find_opt model.vocab text with\n | Some id -> [| id |]\n | None -> (\n match Hashtbl.find_opt model.vocab model.unk_token with\n | Some unk_id -> [| unk_id |]\n | None -> [||])\n\nlet token_to_id model token = Hashtbl.find_opt model.vocab token\nlet id_to_token model id = Hashtbl.find_opt model.vocab_r id\n\nlet get_vocab model =\n Hashtbl.fold (fun token id acc -> (token, id) :: acc) model.vocab []\n\nlet get_vocab_size model = Hashtbl.length model.vocab\n\nlet add_tokens model tokens =\n let start_id = Hashtbl.length model.vocab in\n let count = ref 0 in\n List.iteri\n (fun i token ->\n if not (Hashtbl.mem model.vocab token) then (\n add_token model.vocab model.vocab_r token (start_id + i);\n incr count))\n tokens;\n !count\n\nlet json_obj pairs =\n Jsont.Json.object' (List.map (fun (k, v) -> (Jsont.Json.name k, v)) pairs)\n\nlet json_to_string j =\n match Jsont_bytesrw.encode_string ~format:Jsont.Minify Jsont.json j with\n | Ok s -> s\n | Error e -> failwith e\n\nlet save model ~folder () =\n let vocab_items =\n get_vocab model\n |> List.sort (fun (_, id1) (_, id2) -> compare id1 id2)\n |> List.map (fun (token, id) ->\n json_obj\n [ (\"token\", Jsont.Json.string token); (\"id\", Jsont.Json.int id) ])\n in\n let json =\n json_obj\n [\n (\"type\", Jsont.Json.string \"WordLevel\");\n (\"unk_token\", Jsont.Json.string model.unk_token);\n (\"vocab\", Jsont.Json.list vocab_items);\n ]\n in\n let path = Filename.concat folder \"wordlevel.json\" in\n let oc = open_out path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc (json_to_string json));\n [ \"wordlevel.json\" ]\n\nlet train ~vocab_size ~min_frequency ~show_progress ~special_tokens texts\n existing =\n let _ = show_progress in\n let counts = Hashtbl.create 10000 in\n List.iter\n (fun line ->\n let words = Re.split (Re.compile (Re.rep1 (Re.set \" \\t\\n\\r\"))) line in\n List.iter\n (fun word ->\n if word <> \"\" then\n Hashtbl.replace counts word\n (1 + Option.value ~default:0 (Hashtbl.find_opt counts word)))\n words)\n texts;\n\n let items =\n Hashtbl.fold\n (fun word count acc ->\n if count >= min_frequency then (word, count) :: acc else acc)\n counts []\n |> List.sort (fun (_, c1) (_, c2) -> compare c2 c1)\n in\n let vocab_items = ref [] in\n let idx = ref 0 in\n List.iter\n (fun token ->\n if !idx < vocab_size then (\n vocab_items := (fst token, !idx) :: !vocab_items;\n incr idx))\n items;\n let vocab_items = List.rev !vocab_items in\n\n let specials = List.mapi (fun i token -> (token, i)) special_tokens in\n let vocab = specials @ vocab_items in\n let model =\n match existing with\n | Some model ->\n model.vocab |> Hashtbl.clear;\n model.vocab_r |> Hashtbl.clear;\n List.iter\n (fun (token, id) -> add_token model.vocab model.vocab_r token id)\n vocab;\n model\n | None -> create ~vocab ()\n in\n (model, special_tokens)\n"
},
{
"path": "packages/brot/lib/word_level.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Word-level tokenization model.\n\n {b Internal module.} Direct vocabulary lookup with no subword splitting.\n Each input word is mapped to a single token ID via exact string match. Words\n not in the vocabulary are replaced by [unk_token]. *)\n\ntype t\n(** The type for word-level models. *)\n\n(** {1:creation Creation} *)\n\nval create : ?vocab:(string * int) list -> ?unk_token:string -> unit -> t\n(** [create ?vocab ?unk_token ()] is a word-level model.\n\n - [vocab] is the initial vocabulary as [(token, id)] pairs. Defaults to\n [[]].\n - [unk_token] is the token emitted for unknown words. Defaults to [\"[UNK]\"].\n*)\n\n(** {1:tokenization Tokenization} *)\n\nval tokenize : t -> string -> (int * string * (int * int)) list\n(** [tokenize t s] is [[(id, token, (start, stop))]] for [s]. If [s] is not in\n the vocabulary, [unk_token] is used. If [unk_token] itself is not in the\n vocabulary, the empty list is returned. *)\n\nval tokenize_ids : t -> string -> int array\n(** [tokenize_ids t s] is like {!tokenize} but returns only token IDs. *)\n\n(** {1:vocabulary Vocabulary} *)\n\nval token_to_id : t -> string -> int option\n(** [token_to_id t tok] is the ID of [tok] in the vocabulary. *)\n\nval id_to_token : t -> int -> string option\n(** [id_to_token t id] is the token string for [id]. *)\n\nval get_vocab : t -> (string * int) list\n(** [get_vocab t] is the vocabulary as [(token, id)] pairs. *)\n\nval get_vocab_size : t -> int\n(** [get_vocab_size t] is the number of tokens in the vocabulary. *)\n\nval add_tokens : t -> string list -> int\n(** [add_tokens t toks] adds [toks] to the vocabulary, assigning consecutive IDs\n starting after the current maximum. Returns the number of new tokens\n actually added (duplicates are skipped). Mutates [t]. *)\n\n(** {1:serialization Serialization} *)\n\nval save : t -> folder:string -> unit -> string list\n(** [save t ~folder ()] writes [wordlevel.json] to [folder]. The file contains\n the vocabulary and [unk_token] in JSON format. Returns the list of created\n filenames. *)\n\n(** {1:training Training} *)\n\nval train :\n vocab_size:int ->\n min_frequency:int ->\n show_progress:bool ->\n special_tokens:string list ->\n string list ->\n t option ->\n t * string list\n(** [train ~vocab_size ~min_frequency ~show_progress ~special_tokens texts init]\n learns a vocabulary from [texts] by counting word frequencies.\n\n - [vocab_size] is the target vocabulary size.\n - [min_frequency] is the minimum word frequency to include.\n - [show_progress] enables progress output on [stderr].\n - [special_tokens] are added to the vocabulary first.\n - [init], when provided, seeds the vocabulary from an existing model.\n\n Returns [(model, special_tokens)]. *)\n"
},
{
"path": "packages/brot/lib/wordpiece.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype token = { id : int; value : string; offsets : int * int }\n\n(* Compact trie for zero-allocation longest-prefix matching *)\n\ntype trie = {\n trie_ids : int array;\n child_starts : int array;\n edge_bytes : bytes;\n edge_targets : int array;\n (* Flat 256-element arrays for dense nodes (>16 children) — O(1) lookup *)\n flat_nodes : int array array;\n}\n\nlet build_trie vocab =\n if Hashtbl.length vocab = 0 then\n {\n trie_ids = [||];\n child_starts = [| 0 |];\n edge_bytes = Bytes.empty;\n edge_targets = [||];\n flat_nodes = [||];\n }\n else\n let cap = ref 256 in\n let ids = ref (Array.make !cap (-1)) in\n let ch = ref (Array.init !cap (fun _ -> Hashtbl.create 0)) in\n let n = ref 1 in\n !ch.(0) <- Hashtbl.create 64;\n let grow () =\n let new_cap = !cap * 2 in\n let new_ids = Array.make new_cap (-1) in\n Array.blit !ids 0 new_ids 0 !n;\n ids := new_ids;\n let new_ch =\n Array.init new_cap (fun i ->\n if i < !n then !ch.(i) else Hashtbl.create 0)\n in\n ch := new_ch;\n cap := new_cap\n in\n Hashtbl.iter\n (fun key id ->\n let cur = ref 0 in\n for i = 0 to String.length key - 1 do\n let byte = Char.code (String.unsafe_get key i) in\n let child =\n match Hashtbl.find_opt !ch.(!cur) byte with\n | Some c -> c\n | None ->\n if !n >= !cap then grow ();\n let c = !n in\n incr n;\n !ch.(c) <- Hashtbl.create 4;\n Hashtbl.add !ch.(!cur) byte c;\n c\n in\n cur := child\n done;\n !ids.(!cur) <- id)\n vocab;\n let node_count = !n in\n let trie_ids = Array.init node_count (fun i -> !ids.(i)) in\n let child_starts = Array.make (node_count + 1) 0 in\n let total = ref 0 in\n for i = 0 to node_count - 1 do\n child_starts.(i) <- !total;\n total := !total + Hashtbl.length !ch.(i)\n done;\n child_starts.(node_count) <- !total;\n let edge_bytes = Bytes.create !total in\n let edge_targets = Array.make !total 0 in\n let pos = ref 0 in\n for i = 0 to node_count - 1 do\n Hashtbl.iter\n (fun byte child ->\n Bytes.unsafe_set edge_bytes !pos (Char.unsafe_chr byte);\n edge_targets.(!pos) <- child;\n incr pos)\n !ch.(i)\n done;\n (* Sort each node's children by byte value for binary search *)\n for i = 0 to node_count - 1 do\n let start = child_starts.(i) in\n let stop = child_starts.(i + 1) in\n for j = start + 1 to stop - 1 do\n let kb = Bytes.unsafe_get edge_bytes j in\n let kt = edge_targets.(j) in\n let k = ref (j - 1) in\n while !k >= start && Bytes.unsafe_get edge_bytes !k > kb do\n Bytes.unsafe_set edge_bytes (!k + 1) (Bytes.unsafe_get edge_bytes !k);\n edge_targets.(!k + 1) <- edge_targets.(!k);\n decr k\n done;\n Bytes.unsafe_set edge_bytes (!k + 1) kb;\n edge_targets.(!k + 1) <- kt\n done\n done;\n (* Build flat 256-element arrays for dense nodes (>16 children) *)\n let flat_nodes = Array.make node_count [||] in\n for i = 0 to node_count - 1 do\n let start = child_starts.(i) in\n let count = child_starts.(i + 1) - start in\n if count > 16 then begin\n let flat = Array.make 256 (-1) in\n for j = start to start + count - 1 do\n let b = Char.code (Bytes.unsafe_get edge_bytes j) in\n flat.(b) <- Array.unsafe_get edge_targets j\n done;\n flat_nodes.(i) <- flat\n end\n done;\n { trie_ids; child_starts; edge_bytes; edge_targets; flat_nodes }\n\nlet[@inline] trie_step trie node byte =\n let flat = Array.unsafe_get trie.flat_nodes node in\n if Array.length flat > 0 then Array.unsafe_get flat byte\n else\n let lo = ref (Array.unsafe_get trie.child_starts node) in\n let hi = ref (Array.unsafe_get trie.child_starts (node + 1) - 1) in\n let result = ref (-1) in\n while !lo <= !hi do\n let mid = !lo + ((!hi - !lo) asr 1) in\n let mid_byte = Char.code (Bytes.unsafe_get trie.edge_bytes mid) in\n if mid_byte = byte then (\n result := Array.unsafe_get trie.edge_targets mid;\n lo := !hi + 1)\n else if mid_byte < byte then lo := mid + 1\n else hi := mid - 1\n done;\n !result\n\nlet trie_longest_match trie sequence ~start ~prefix ~prefix_len =\n if Array.length trie.trie_ids = 0 then None\n else\n let seq_len = String.length sequence in\n let last_id = ref (-1) in\n let last_end = ref start in\n let current = ref 0 in\n let stopped = ref false in\n let i = ref 0 in\n while !i < prefix_len && not !stopped do\n let child =\n trie_step trie !current (Char.code (String.unsafe_get prefix !i))\n in\n if child < 0 then stopped := true\n else (\n current := child;\n incr i)\n done;\n (if not !stopped then\n let j = ref start in\n while !j < seq_len && not !stopped do\n let child =\n trie_step trie !current (Char.code (String.unsafe_get sequence !j))\n in\n if child < 0 then stopped := true\n else (\n current := child;\n incr j;\n let tid = Array.unsafe_get trie.trie_ids child in\n if tid >= 0 then (\n last_id := tid;\n last_end := !j))\n done);\n if !last_id >= 0 then Some (!last_id, !last_end) else None\n\n(* Model type *)\n\ntype t = {\n vocab : (string, int) Hashtbl.t;\n vocab_r : string array;\n trie : trie;\n unk_token : string;\n continuing_subword_prefix : string;\n max_input_chars_per_word : int;\n}\n\nlet create ~vocab ?(unk_token = \"[UNK]\") ?(continuing_subword_prefix = \"##\")\n ?(max_input_chars_per_word = 100) () =\n let max_id = Hashtbl.fold (fun _ id acc -> max id acc) vocab (-1) in\n let vocab_r = Array.make (max_id + 1) \"\" in\n Hashtbl.iter (fun k v -> Array.unsafe_set vocab_r v k) vocab;\n if Hashtbl.length vocab > 0 && not (Hashtbl.mem vocab unk_token) then\n invalid_arg \"Wordpiece.create: unk_token not in vocab\";\n let trie = build_trie vocab in\n {\n vocab;\n vocab_r;\n trie;\n unk_token;\n continuing_subword_prefix;\n max_input_chars_per_word;\n }\n\nlet read_file ~vocab_file =\n let vocab = Hashtbl.create 10000 in\n let ic = open_in vocab_file in\n Fun.protect\n ~finally:(fun () -> close_in ic)\n (fun () ->\n let index = ref 0 in\n (try\n while true do\n let line = input_line ic in\n let token = String.trim line in\n if token <> \"\" then (\n Hashtbl.add vocab token !index;\n incr index)\n done\n with End_of_file -> ());\n vocab)\n\nlet from_file ~vocab_file =\n let vocab = read_file ~vocab_file in\n create ~vocab ()\n\nlet count_chars s =\n let len = String.length s in\n let n = ref 0 in\n for i = 0 to len - 1 do\n if Char.code (String.unsafe_get s i) land 0xC0 <> 0x80 then incr n\n done;\n !n\n\nlet tokenize model sequence =\n if Hashtbl.length model.vocab = 0 then []\n else\n let seq_len = String.length sequence in\n if count_chars sequence > model.max_input_chars_per_word then\n let id = Hashtbl.find model.vocab model.unk_token in\n [ { id; value = model.unk_token; offsets = (0, seq_len) } ]\n else\n let prefix = model.continuing_subword_prefix in\n let prefix_len = String.length prefix in\n let rec greedy start acc =\n if start >= seq_len then List.rev acc\n else\n let p = if start > 0 then prefix else \"\" in\n let pl = if start > 0 then prefix_len else 0 in\n match\n trie_longest_match model.trie sequence ~start ~prefix:p\n ~prefix_len:pl\n with\n | Some (id, end_byte) ->\n let value = Array.unsafe_get model.vocab_r id in\n greedy end_byte ({ id; value; offsets = (start, end_byte) } :: acc)\n | None ->\n let id = Hashtbl.find model.vocab model.unk_token in\n [ { id; value = model.unk_token; offsets = (0, seq_len) } ]\n in\n greedy 0 []\n\nlet tokenize_ids model sequence =\n if Hashtbl.length model.vocab = 0 then [||]\n else\n let seq_len = String.length sequence in\n if count_chars sequence > model.max_input_chars_per_word then\n let id = Hashtbl.find model.vocab model.unk_token in\n [| id |]\n else\n let prefix = model.continuing_subword_prefix in\n let prefix_len = String.length prefix in\n let ids = ref [] in\n let n = ref 0 in\n let rec greedy start =\n if start >= seq_len then ()\n else\n let p = if start > 0 then prefix else \"\" in\n let pl = if start > 0 then prefix_len else 0 in\n match\n trie_longest_match model.trie sequence ~start ~prefix:p\n ~prefix_len:pl\n with\n | Some (id, end_byte) ->\n ids := id :: !ids;\n incr n;\n greedy end_byte\n | None ->\n let unk_id = Hashtbl.find model.vocab model.unk_token in\n ids := [ unk_id ];\n n := 1\n in\n greedy 0;\n let result = Array.make !n 0 in\n List.iteri (fun i id -> result.(!n - 1 - i) <- id) !ids;\n result\n\nlet tokenize_spans_encoding model pre_tokens ~type_id =\n if Hashtbl.length model.vocab = 0 then Encoding.empty\n else\n let trie = model.trie in\n let prefix = model.continuing_subword_prefix in\n let prefix_len = String.length prefix in\n let unk_id = Hashtbl.find model.vocab model.unk_token in\n let max_chars = model.max_input_chars_per_word in\n let vocab_r = model.vocab_r in\n let unk_token_str = model.unk_token in\n (* Single pass: convert pre_tokens to array for direct access (no closure),\n tokenize all fragments and fill growable output arrays directly. *)\n let pre_arr = Array.of_list pre_tokens in\n let n_pre = Array.length pre_arr in\n let cap = ref (max 16 (n_pre * 2)) in\n let ids = ref (Array.make !cap 0) in\n let token_strs = ref (Array.make !cap \"\") in\n let offsets_arr = ref (Array.make !cap (0, 0)) in\n let n = ref 0 in\n let grow () =\n let new_cap = !cap * 2 in\n let new_ids = Array.make new_cap 0 in\n Array.blit !ids 0 new_ids 0 !n;\n ids := new_ids;\n let new_strs = Array.make new_cap \"\" in\n Array.blit !token_strs 0 new_strs 0 !n;\n token_strs := new_strs;\n let new_off = Array.make new_cap (0, 0) in\n Array.blit !offsets_arr 0 new_off 0 !n;\n offsets_arr := new_off;\n cap := new_cap\n in\n (* Hoisted mutable state for trie matching — allocated once *)\n let current = ref 0 in\n let stopped = ref false in\n let last_id = ref (-1) in\n let last_end = ref 0 in\n let pos = ref 0 in\n let is_unk = ref false in\n let char_count = ref 0 in\n let i_ref = ref 0 in\n let j_ref = ref 0 in\n for frag_idx = 0 to n_pre - 1 do\n let fragment, _ = Array.unsafe_get pre_arr frag_idx in\n let seq_len = String.length fragment in\n char_count := 0;\n for k = 0 to seq_len - 1 do\n if Char.code (String.unsafe_get fragment k) land 0xC0 <> 0x80 then\n incr char_count\n done;\n if !char_count > max_chars then begin\n if !n >= !cap then grow ();\n Array.unsafe_set !ids !n unk_id;\n Array.unsafe_set !token_strs !n unk_token_str;\n Array.unsafe_set !offsets_arr !n (0, seq_len);\n incr n\n end\n else begin\n pos := 0;\n is_unk := false;\n let start_n = !n in\n while !pos < seq_len && not !is_unk do\n let match_start = !pos in\n current := 0;\n stopped := false;\n last_id := -1;\n last_end := !pos;\n if !pos > 0 then begin\n i_ref := 0;\n while !i_ref < prefix_len && not !stopped do\n let child =\n trie_step trie !current\n (Char.code (String.unsafe_get prefix !i_ref))\n in\n if child < 0 then stopped := true\n else begin\n current := child;\n incr i_ref\n end\n done\n end;\n if not !stopped then begin\n j_ref := !pos;\n while !j_ref < seq_len && not !stopped do\n let child =\n trie_step trie !current\n (Char.code (String.unsafe_get fragment !j_ref))\n in\n if child < 0 then stopped := true\n else begin\n current := child;\n incr j_ref;\n let tid = Array.unsafe_get trie.trie_ids child in\n if tid >= 0 then begin\n last_id := tid;\n last_end := !j_ref\n end\n end\n done\n end;\n if !last_id >= 0 then begin\n if !n >= !cap then grow ();\n Array.unsafe_set !ids !n !last_id;\n Array.unsafe_set !token_strs !n (Array.unsafe_get vocab_r !last_id);\n Array.unsafe_set !offsets_arr !n (match_start, !last_end);\n incr n;\n pos := !last_end\n end\n else is_unk := true\n done;\n if !is_unk then begin\n n := start_n;\n if !n >= !cap then grow ();\n Array.unsafe_set !ids !n unk_id;\n Array.unsafe_set !token_strs !n unk_token_str;\n Array.unsafe_set !offsets_arr !n (0, seq_len);\n n := start_n + 1\n end\n end\n done;\n let total = !n in\n if total = 0 then Encoding.empty\n else\n let final_ids = if total = !cap then !ids else Array.sub !ids 0 total in\n let final_strs =\n if total = !cap then !token_strs else Array.sub !token_strs 0 total\n in\n let final_off =\n if total = !cap then !offsets_arr else Array.sub !offsets_arr 0 total\n in\n Encoding.create ~ids:final_ids ~type_ids:(Array.make total type_id)\n ~tokens:final_strs ~words:(Array.make total None) ~offsets:final_off\n ~special_tokens_mask:(Array.make total 0)\n ~attention_mask:(Array.make total 1) ()\n\nlet token_to_id model token = Hashtbl.find_opt model.vocab token\n\nlet id_to_token model id =\n if id >= 0 && id < Array.length model.vocab_r then\n Some (Array.unsafe_get model.vocab_r id)\n else None\n\nlet get_vocab model = Hashtbl.fold (fun k v acc -> (k, v) :: acc) model.vocab []\nlet get_vocab_size model = Hashtbl.length model.vocab\nlet get_unk_token model = model.unk_token\nlet get_continuing_subword_prefix model = model.continuing_subword_prefix\n\nlet save model ~path ?name () =\n let vocab_file =\n match name with\n | Some n -> Filename.concat path (n ^ \"-vocab.txt\")\n | None -> Filename.concat path \"vocab.txt\"\n in\n let vocab_list =\n Hashtbl.fold (fun k v acc -> (v, k) :: acc) model.vocab []\n |> List.sort compare\n |> List.map (fun (_, k) -> k)\n in\n let oc = open_out vocab_file in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () ->\n List.iter\n (fun token ->\n output_string oc token;\n output_char oc '\\n')\n vocab_list);\n vocab_file\n\nlet from_bpe bpe =\n let vocab = Hashtbl.create (Bpe.get_vocab_size bpe) in\n List.iter (fun (k, id) -> Hashtbl.add vocab k id) (Bpe.get_vocab bpe);\n let unk_token =\n match Bpe.get_unk_token bpe with Some u -> u | None -> \"[UNK]\"\n in\n if not (Hashtbl.mem vocab unk_token) then begin\n let max_id = Hashtbl.fold (fun _ id acc -> max id acc) vocab (-1) in\n Hashtbl.add vocab unk_token (max_id + 1)\n end;\n let continuing_subword_prefix =\n match Bpe.get_continuing_subword_prefix bpe with\n | Some p -> p\n | None -> \"##\"\n in\n create ~vocab ~unk_token ~continuing_subword_prefix ()\n\n(* Trainer *)\n\nlet train ~min_frequency ~vocab_size ~show_progress ~special_tokens\n ~limit_alphabet ~initial_alphabet ~continuing_subword_prefix\n ~end_of_word_suffix texts existing =\n let _ = existing in\n (* WordPiece training uses BPE algorithm internally *)\n let bpe_trained, result_tokens =\n Bpe.train ~min_frequency ~vocab_size ~show_progress ~special_tokens\n ~limit_alphabet ~initial_alphabet\n ~continuing_subword_prefix:(Some continuing_subword_prefix)\n ~end_of_word_suffix ~max_token_length:None texts None\n in\n let wordpiece_model = from_bpe bpe_trained in\n (wordpiece_model, result_tokens)\n"
},
{
"path": "packages/brot/lib/wordpiece.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** WordPiece tokenization model.\n\n {b Internal module.} Greedy longest-match-first subword decomposition\n against a fixed vocabulary. Used by BERT, DistilBERT, and Electra.\n\n A word is decomposed left-to-right: at each position the longest vocabulary\n match is consumed. Continuation pieces are prefixed with\n {!get_continuing_subword_prefix} (typically [\"##\"]). If no subword is found\n at any position the {e entire} word falls back to {!get_unk_token}.\n\n Vocabulary lookup uses a hybrid trie: dense nodes (more than 16 children)\n use a 256-element flat array for O(1) byte dispatch, sparse nodes use binary\n search on sorted edges. *)\n\ntype t\n(** The type for WordPiece models. *)\n\n(** {1:creation Creation} *)\n\nval create :\n vocab:(string, int) Hashtbl.t ->\n ?unk_token:string ->\n ?continuing_subword_prefix:string ->\n ?max_input_chars_per_word:int ->\n unit ->\n t\n(** [create ~vocab ()] is a WordPiece model backed by [vocab].\n\n - [unk_token] is the token emitted for words that cannot be decomposed.\n Defaults to [\"[UNK]\"].\n - [continuing_subword_prefix] is prepended to non-initial subwords. Defaults\n to [\"##\"].\n - [max_input_chars_per_word] is the UTF-8 character count above which a word\n is replaced by [unk_token] without attempting decomposition. Defaults to\n [100].\n\n Raises [Invalid_argument] if [unk_token] is not in [vocab]. *)\n\nval from_file : vocab_file:string -> t\n(** [from_file ~vocab_file] loads a model from a BERT-style [vocab.txt] file\n (one token per line, ID equals line number). Uses BERT defaults:\n [unk_token = \"[UNK]\"], [continuing_subword_prefix = \"##\"],\n [max_input_chars_per_word = 100]. *)\n\n(** {1:tokenization Tokenization} *)\n\ntype token = { id : int; value : string; offsets : int * int }\n(** The type for tokens. [id] is the vocabulary index, [value] the string\n content, and [offsets] the [(start, stop)] byte span in the source text. *)\n\nval tokenize : t -> string -> token list\n(** [tokenize t s] is the WordPiece decomposition of [s].\n\n If [s] exceeds {!create}'s [max_input_chars_per_word] (in UTF-8 characters),\n a single [unk_token] token spanning the whole input is returned. If\n decomposition fails at any position, the result is likewise a single\n [unk_token]. *)\n\nval tokenize_ids : t -> string -> int array\n(** [tokenize_ids t s] is like {!tokenize} but returns only token IDs. *)\n\nval tokenize_spans_encoding :\n t -> (string * (int * int)) list -> type_id:int -> Encoding.t\n(** [tokenize_spans_encoding t spans ~type_id] tokenizes all [spans] and builds\n an {!Encoding.t} directly. Each element of [spans] is\n [(fragment, (start, stop))] where offsets are byte positions in the original\n text.\n\n This is a single-pass variant that avoids intermediate list and record\n allocation: mutable refs are hoisted, growable arrays are filled in place,\n and trie matching is inlined. *)\n\n(** {1:vocabulary Vocabulary} *)\n\nval token_to_id : t -> string -> int option\n(** [token_to_id t tok] is the ID of [tok] in the vocabulary. *)\n\nval id_to_token : t -> int -> string option\n(** [id_to_token t id] is the token string for [id]. *)\n\nval get_vocab : t -> (string * int) list\n(** [get_vocab t] is the vocabulary as [(token, id)] pairs. *)\n\nval get_vocab_size : t -> int\n(** [get_vocab_size t] is the number of tokens in the vocabulary. *)\n\nval get_unk_token : t -> string\n(** [get_unk_token t] is the unknown token string. *)\n\nval get_continuing_subword_prefix : t -> string\n(** [get_continuing_subword_prefix t] is the subword continuation prefix (e.g.\n [\"##\"]). *)\n\n(** {1:serialization Serialization} *)\n\nval save : t -> path:string -> ?name:string -> unit -> string\n(** [save t ~path ()] writes the vocabulary as a plain-text [vocab.txt] file\n (one token per line) to [path]. If [name] is given the file is named\n [{name}-vocab.txt]. Returns the filepath written. *)\n\n(** {1:training Training} *)\n\nval train :\n min_frequency:int ->\n vocab_size:int ->\n show_progress:bool ->\n special_tokens:string list ->\n limit_alphabet:int option ->\n initial_alphabet:char list ->\n continuing_subword_prefix:string ->\n end_of_word_suffix:string option ->\n string list ->\n t option ->\n t * string list\n(** [train ~min_frequency ~vocab_size ~show_progress ~special_tokens\n ~limit_alphabet ~initial_alphabet ~continuing_subword_prefix\n ~end_of_word_suffix texts init] learns a WordPiece vocabulary from [texts]\n using BPE merge training internally.\n\n - [min_frequency] is the minimum pair frequency to merge.\n - [vocab_size] is the target vocabulary size.\n - [show_progress] enables progress output on [stderr].\n - [special_tokens] are added to the vocabulary first.\n - [limit_alphabet] caps the number of distinct initial characters kept.\n - [initial_alphabet] seeds the character set.\n - [continuing_subword_prefix] is set on the resulting model.\n - [end_of_word_suffix] appended to final subwords if given.\n - [init], when provided, seeds the vocabulary from an existing model.\n\n Returns [(model, special_tokens)]. *)\n"
},
{
"path": "packages/brot/test/dune",
"content": "(data_only_dirs fixtures scripts)\n\n(tests\n (names\n test_tokenization\n test_vocab\n test_encoding\n test_unicode\n test_bpe\n test_wordpiece\n test_hf_tokenizers\n test_processors\n test_pretokenizers)\n (package brot)\n (libraries brot windtrap unix jsont))\n"
},
{
"path": "packages/brot/test/fixtures/.gitignore",
"content": "hf/\n"
},
{
"path": "packages/brot/test/scripts/download_hf_tokenizers.py",
"content": "#!/usr/bin/env python3\n\"\"\"\nDownload selected HuggingFace tokenizer JSON files into brot/test/fixtures/hf.\n\nRun this script whenever you need to refresh the fixtures:\n\n python3 brot/test/scripts/download_hf_tokenizers.py\n\nThe files are ignored by git, so each developer/machine maintains its own cache.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport hashlib\nimport json\nimport sys\nimport urllib.request\nfrom pathlib import Path\nfrom typing import Iterable, Tuple\n\n\nFIXTURES: Iterable[Tuple[str, str]] = (\n (\n \"bert-base-uncased\",\n \"https://huggingface.co/bert-base-uncased/resolve/main/tokenizer.json?download=1\",\n ),\n (\n \"gpt2\",\n \"https://huggingface.co/gpt2/resolve/main/tokenizer.json?download=1\",\n ),\n (\n \"roberta-base\",\n \"https://huggingface.co/roberta-base/resolve/main/tokenizer.json?download=1\",\n ),\n)\n\n\ndef download(url: str, dest: Path) -> None:\n dest.parent.mkdir(parents=True, exist_ok=True)\n tmp_path = dest.with_suffix(\".tmp\")\n\n print(f\"→ downloading {url}…\")\n with urllib.request.urlopen(url) as response, open(tmp_path, \"wb\") as out:\n while True:\n chunk = response.read(1024 * 64)\n if not chunk:\n break\n out.write(chunk)\n tmp_path.replace(dest)\n\n\ndef sha256(path: Path) -> str:\n h = hashlib.sha256()\n with path.open(\"rb\") as fh:\n for chunk in iter(lambda: fh.read(1024 * 64), b\"\"):\n h.update(chunk)\n return h.hexdigest()\n\n\ndef summarize(path: Path) -> None:\n try:\n with path.open(\"r\", encoding=\"utf-8\") as fh:\n metadata = json.load(fh)\n model_type = metadata.get(\"model\", {}).get(\"type\", \"\")\n size = path.stat().st_size\n digest = sha256(path)[:12]\n print(f\" saved {path} ({size} bytes, model={model_type}, sha256={digest})\")\n except Exception as exc: # pylint: disable=broad-except\n print(f\" warning: failed to inspect {path}: {exc}\")\n\n\ndef main() -> int:\n test_root = Path(__file__).resolve().parents[1]\n fixtures_dir = test_root / \"fixtures\" / \"hf\"\n fixtures_dir.mkdir(parents=True, exist_ok=True)\n\n for model, url in FIXTURES:\n target = fixtures_dir / model / \"tokenizer.json\"\n if target.exists():\n print(f\"✓ {model} already present at {target}\")\n continue\n print(f\"Downloading {model} tokenizer…\")\n try:\n download(url, target)\n summarize(target)\n except Exception as exc: # pylint: disable=broad-except\n print(f\" failed to download {model}: {exc}\", file=sys.stderr)\n if target.exists():\n target.unlink()\n return 1\n\n print(\"All fixtures downloaded.\")\n return 0\n\n\nif __name__ == \"__main__\":\n sys.exit(main())\n"
},
{
"path": "packages/brot/test/test_bpe.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nopen Brot\n\nlet test_bpe_basic () =\n (* Create a simple vocabulary and merges *)\n let vocab =\n [\n (\"h\", 0);\n (\"e\", 1);\n (\"l\", 2);\n (\"o\", 3);\n (\"ll\", 4);\n (\"he\", 5);\n (\"llo\", 6);\n (\"hello\", 7);\n ]\n in\n\n let merges =\n [\n (\"l\", \"l\");\n (* rank 0: Merge 'l' + 'l' -> 'll' *)\n (\"ll\", \"o\");\n (* rank 1: Merge 'll' + 'o' -> 'llo' *)\n (\"he\", \"llo\");\n (* rank 2: Merge 'he' + 'llo' -> 'hello' *)\n ]\n in\n\n let tokenizer = bpe ~vocab ~merges ~unk_token:\"\" () in\n\n let encoding = encode tokenizer \"hello\" in\n let tokens = Encoding.tokens encoding |> Array.to_list in\n\n Printf.printf \"Tokenized 'hello': \";\n List.iter (Printf.printf \"%s \") tokens;\n Printf.printf \"\\n\";\n\n equal ~msg:\"vocabulary size\" int 8 (vocab_size tokenizer)\n\nlet test_bpe_builder () =\n let vocab = [ (\"a\", 0); (\"b\", 1); (\"ab\", 2) ] in\n let merges = [ (\"a\", \"b\") ] in\n\n let tokenizer = bpe ~vocab ~merges ~cache_capacity:50 () in\n\n let encoding = encode tokenizer \"ab\" in\n let tokens = Encoding.tokens encoding in\n equal ~msg:\"single token for 'ab'\" int 1 (Array.length tokens)\n\nlet test_bpe_save_load () =\n let vocab = [ (\"t\", 0); (\"e\", 1); (\"s\", 2); (\"test\", 3) ] in\n let merges = [] in\n (* No merges for simplicity *)\n\n let tokenizer = bpe ~vocab ~merges () in\n\n (* Save the model *)\n let temp_dir = Filename.temp_dir \"bpe_test\" \"\" in\n let files = save_model_files tokenizer ~folder:temp_dir () in\n\n (* Load the model *)\n let vocab_file = List.find (fun f -> Filename.check_suffix f \".json\") files in\n let merges_file = List.find (fun f -> Filename.check_suffix f \".txt\") files in\n let loaded_tokenizer =\n from_model_file ~vocab:vocab_file ~merges:merges_file ()\n in\n\n (* Test that loaded tokenizer works the same *)\n let original_tokens = encode tokenizer \"test\" |> Encoding.tokens in\n let loaded_tokens = encode loaded_tokenizer \"test\" |> Encoding.tokens in\n\n equal ~msg:\"same number of tokens\" int\n (Array.length original_tokens)\n (Array.length loaded_tokens);\n\n (* Clean up *)\n List.iter Sys.remove files;\n Unix.rmdir temp_dir\n\nlet test_tokenizer_integration () =\n (* Create a BPE tokenizer using the high-level API *)\n let vocab =\n [\n (\"h\", 0); (\"e\", 1); (\"l\", 2); (\"o\", 3); (\"he\", 4); (\"llo\", 5); (\"hello\", 6);\n ]\n in\n let merges = [ (\"h\", \"e\"); (\"he\", \"llo\") ] in\n let tokenizer = bpe ~vocab ~merges () in\n\n (* Test encoding *)\n let tokens = encode tokenizer \"hello\" |> Encoding.tokens |> Array.to_list in\n\n Printf.printf \"bpe result: \";\n List.iter (Printf.printf \"%s \") tokens;\n Printf.printf \"\\n\";\n\n equal ~msg:\"tokenizer produces output\" bool true (List.length tokens > 0)\n\nlet () =\n run \"BPE tests\"\n [\n group \"basic\"\n [\n test \"basic tokenization\" test_bpe_basic;\n test \"builder pattern\" test_bpe_builder;\n test \"save and load\" test_bpe_save_load;\n test \"tokenizer integration\" test_tokenizer_integration;\n ];\n ]\n"
},
{
"path": "packages/brot/test/test_encoding.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nopen Brot\n\nlet make_word_tokenizer ?(specials = []) () =\n word_level ~pre:(Pre_tokenizer.whitespace ()) ~specials ()\n\nlet test_encode_simple () =\n let tokenizer = add_tokens (make_word_tokenizer ()) [ \"hello\"; \"world\" ] in\n let ids = encode tokenizer \"hello world hello\" |> Encoding.ids in\n equal ~msg:\"encoded length\" int 3 (Array.length ids);\n equal ~msg:\"repeated token same id\" bool true (ids.(0) = ids.(2))\n\nlet test_encode_with_vocab () =\n let tokenizer = add_tokens (make_word_tokenizer ()) [ \"hello\"; \"world\" ] in\n let ids = encode tokenizer \"hello world\" |> Encoding.ids |> Array.to_list in\n equal ~msg:\"encoded with vocab\" (list int) [ 0; 1 ] ids\n\nlet test_encode_unknown_tokens () =\n let tokenizer =\n add_tokens\n (make_word_tokenizer ~specials:[ special \"\" ] ())\n [ \"hello\" ]\n in\n let ids =\n encode tokenizer \"hello unknown world\" |> Encoding.ids |> Array.to_list\n in\n equal ~msg:\"encoded something\" bool true (List.length ids > 0)\n\nlet test_encode_empty () =\n let tokenizer = make_word_tokenizer () in\n let ids = encode tokenizer \"\" |> Encoding.ids |> Array.to_list in\n equal ~msg:\"encode empty\" (list int) [] ids\n\nlet test_encode_batch_simple () =\n let tokenizer =\n add_tokens (make_word_tokenizer ()) [ \"hello\"; \"world\"; \"hi\"; \"there\" ]\n in\n let encodings = encode_batch tokenizer [ \"hello world\"; \"hi there\" ] in\n equal ~msg:\"batch size\" int 2 (List.length encodings);\n let first = List.hd encodings in\n equal ~msg:\"first encoding has ids\" bool true\n (Array.length (Encoding.ids first) > 0)\n\nlet test_encode_batch_with_padding () =\n let tokenizer =\n add_tokens\n (make_word_tokenizer ~specials:[ special \"\" ] ())\n [ \"hello\"; \"world\"; \"hi\"; \"there\" ]\n in\n let padding =\n {\n length = `Fixed 5;\n direction = `Right;\n pad_id = None;\n pad_type_id = None;\n pad_token = Some \"\";\n }\n in\n let encodings = encode_batch tokenizer ~padding [ \"hello\"; \"hi there\" ] in\n let first = Encoding.ids (List.nth encodings 0) in\n let second = Encoding.ids (List.nth encodings 1) in\n equal ~msg:\"first padded length\" int 5 (Array.length first);\n equal ~msg:\"second padded length\" int 5 (Array.length second)\n\nlet test_encode_batch_empty () =\n let tokenizer = make_word_tokenizer () in\n let encodings = encode_batch tokenizer [] in\n equal ~msg:\"empty batch\" int 0 (List.length encodings)\n\nlet test_decode_simple () =\n let tokenizer = add_tokens (make_word_tokenizer ()) [ \"hello\"; \"world\" ] in\n let decoded = decode tokenizer [| 0; 1 |] in\n equal ~msg:\"decoded text\" string \"hello world\" decoded\n\nlet test_decode_with_special () =\n let tokenizer =\n add_tokens\n (make_word_tokenizer ~specials:[ special \"\"; special \"\" ] ())\n [ \"hello\" ]\n in\n (* =0, =1, hello=2 *)\n let decoded = decode tokenizer [| 0; 2; 1 |] in\n equal ~msg:\"decoded with special\" string \" hello \" decoded\n\nlet test_decode_skip_special () =\n let tokenizer =\n add_tokens\n (make_word_tokenizer ~specials:[ special \"\"; special \"\" ] ())\n [ \"hello\" ]\n in\n let decoded = decode ~skip_special_tokens:true tokenizer [| 0; 2; 1 |] in\n equal ~msg:\"decoded without special\" string \"hello\" decoded\n\nlet test_decode_batch () =\n let tokenizer =\n add_tokens (make_word_tokenizer ()) [ \"hello\"; \"world\"; \"hi\"; \"there\" ]\n in\n let decoded = decode_batch tokenizer [ [| 0; 1 |]; [| 2; 3 |] ] in\n equal ~msg:\"decoded count\" int 2 (List.length decoded);\n equal ~msg:\"first decoded\" string \"hello world\" (List.nth decoded 0);\n equal ~msg:\"second decoded\" string \"hi there\" (List.nth decoded 1)\n\nlet test_chars_model () =\n let tokenizer = chars () in\n let ids = encode tokenizer \"abc\" |> Encoding.ids |> Array.to_list in\n equal ~msg:\"char ids\" (list int) [ 97; 98; 99 ] ids\n\nlet suite =\n [\n test \"encode simple\" test_encode_simple;\n test \"encode with vocab\" test_encode_with_vocab;\n test \"encode unknown tokens\" test_encode_unknown_tokens;\n test \"encode empty\" test_encode_empty;\n test \"batch simple\" test_encode_batch_simple;\n test \"batch with padding\" test_encode_batch_with_padding;\n test \"batch empty request\" test_encode_batch_empty;\n test \"decode simple\" test_decode_simple;\n test \"decode with special\" test_decode_with_special;\n test \"decode skip special\" test_decode_skip_special;\n test \"decode batch\" test_decode_batch;\n test \"chars model\" test_chars_model;\n ]\n\nlet () = run \"Encoding tests\" [ group \"encoding\" suite ]\n"
},
{
"path": "packages/brot/test/test_hf_tokenizers.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Brot\nopen Windtrap\n\nlet candidate_roots () =\n match Sys.getenv_opt \"DUNE_SOURCEROOT\" with\n | Some root -> [ root; Sys.getcwd () ]\n | None -> [ Sys.getcwd () ]\n\nlet locate_fixture model =\n let relative =\n Filename.concat \"brot/test/fixtures/hf\"\n (Filename.concat model \"tokenizer.json\")\n in\n let rec search = function\n | [] -> None\n | root :: rest ->\n let path = Filename.concat root relative in\n if Sys.file_exists path then Some path else search rest\n in\n search (candidate_roots ())\n\nlet with_hf_tokenizer model f =\n match locate_fixture model with\n | None -> skip ()\n | Some path -> (\n match from_file path with\n | Ok tok -> f tok\n | Error msg -> failf \"Failed to load tokenizer %s: %s\" model msg)\n\nlet test_bert_base_uncased () =\n with_hf_tokenizer \"bert-base-uncased\" (fun tok ->\n let encoding = encode tok \"Hello world!\" in\n let tokens = Encoding.tokens encoding |> Array.to_list in\n equal ~msg:\"token sequence\" (list string)\n [ \"[CLS]\"; \"hello\"; \"world\"; \"!\"; \"[SEP]\" ]\n tokens;\n let type_ids = Encoding.type_ids encoding |> Array.to_list in\n equal ~msg:\"type ids\" (list int) [ 0; 0; 0; 0; 0 ] type_ids;\n equal ~msg:\"has [MASK]\" bool true\n (Option.is_some (token_to_id tok \"[MASK]\")))\n\nlet test_gpt2_small () =\n with_hf_tokenizer \"gpt2\" (fun tok ->\n let encoding = encode tok \"Hello world\" in\n let ids = Encoding.ids encoding |> Array.to_list in\n equal ~msg:\"ids\" (list int) [ 15496; 995 ] ids;\n let roundtrip =\n decode tok (Array.of_list ids) ~skip_special_tokens:true\n in\n equal ~msg:\"decode\" string \"Hello world\" roundtrip)\n\nlet test_roberta_base () =\n with_hf_tokenizer \"roberta-base\" (fun tok ->\n let encoding = encode tok \"A quick test\" in\n let tokens = Encoding.tokens encoding |> Array.to_list in\n equal ~msg:\"tokens\" (list string)\n [ \"\"; \"A\"; \"Ġquick\"; \"Ġtest\"; \" \" ]\n tokens;\n let attention = Encoding.attention_mask encoding |> Array.to_list in\n equal ~msg:\"attention mask\" (list int) [ 1; 1; 1; 1; 1 ] attention)\n\nlet () =\n run \"HF tokenizers\"\n [\n group \"bert-base-uncased\" [ test \"encode\" test_bert_base_uncased ];\n group \"gpt2\" [ test \"encode\" test_gpt2_small ];\n group \"roberta-base\" [ test \"encode\" test_roberta_base ];\n ]\n"
},
{
"path": "packages/brot/test/test_pretokenizers.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Pre = Brot.Pre_tokenizer\n\nlet check_tokenization name input expected =\n equal ~msg:name (list (pair string (pair int int))) expected input\n\nlet check_strings name input expected =\n equal ~msg:name (list string) expected (List.map fst input)\n\nlet test_byte_level_basic () =\n let tokenizer = Pre.byte_level ~add_prefix_space:false ~use_regex:true () in\n\n (* Test basic tokenization *)\n let test_case text expected_pieces expected_offsets =\n let result = Pre.pre_tokenize tokenizer text in\n let offsets = List.map snd result in\n check_strings\n (Printf.sprintf \"ByteLevel pieces for %S\" text)\n result expected_pieces;\n equal\n ~msg:(Printf.sprintf \"ByteLevel offsets for %S\" text)\n (list (pair int int))\n expected_offsets offsets\n in\n\n (* Basic words *)\n test_case \"Hello\" [ \"Hello\" ] [ (0, 5) ];\n test_case \"hello\" [ \"hello\" ] [ (0, 5) ];\n test_case \"HELLO\" [ \"HELLO\" ] [ (0, 5) ];\n\n (* Words with spaces - space becomes Ġ (0xC4 0xA0) *)\n test_case \"Hello world\" [ \"Hello\"; \"\\196\\160world\" ] [ (0, 5); (5, 11) ];\n test_case \"Hello world\"\n [ \"Hello\"; \"\\196\\160\"; \"\\196\\160world\" ]\n [ (0, 5); (5, 6); (6, 12) ];\n\n (* Leading/trailing spaces *)\n test_case \" hello\" [ \"\\196\\160hello\" ] [ (0, 6) ];\n test_case \"hello \" [ \"hello\"; \"\\196\\160\" ] [ (0, 5); (5, 6) ];\n (* Note: Python produces ['Ġ', 'Ġhello', 'ĠĠ'] for \" hello \" *)\n test_case \" hello \"\n [ \"\\196\\160\"; \"\\196\\160hello\"; \"\\196\\160\\196\\160\" ]\n [ (0, 1); (1, 7); (7, 9) ];\n\n (* Contractions - should be kept as separate pieces *)\n test_case \"'s\" [ \"'s\" ] [ (0, 2) ];\n test_case \"'t\" [ \"'t\" ] [ (0, 2) ];\n test_case \"'re\" [ \"'re\" ] [ (0, 3) ];\n test_case \"'ve\" [ \"'ve\" ] [ (0, 3) ];\n test_case \"'m\" [ \"'m\" ] [ (0, 2) ];\n test_case \"'ll\" [ \"'ll\" ] [ (0, 3) ];\n test_case \"'d\" [ \"'d\" ] [ (0, 2) ];\n\n (* Words with contractions *)\n test_case \"don't\" [ \"don\"; \"'t\" ] [ (0, 3); (3, 5) ];\n test_case \"it's\" [ \"it\"; \"'s\" ] [ (0, 2); (2, 4) ];\n test_case \"we're\" [ \"we\"; \"'re\" ] [ (0, 2); (2, 5) ];\n test_case \"I'll\" [ \"I\"; \"'ll\" ] [ (0, 1); (1, 4) ];\n test_case \"OpenAI's\" [ \"OpenAI\"; \"'s\" ] [ (0, 6); (6, 8) ]\n\nlet test_byte_level_prefix_space () =\n (* Test with add_prefix_space=true *)\n let tokenizer = Pre.byte_level ~add_prefix_space:true ~use_regex:true () in\n\n let test_case text expected_pieces =\n let result = Pre.pre_tokenize tokenizer text in\n check_strings\n (Printf.sprintf \"ByteLevel with prefix for %S\" text)\n result expected_pieces\n in\n\n (* Should add space prefix when text doesn't start with space *)\n test_case \"hello\" [ \"\\196\\160hello\" ];\n test_case \"Hello world\" [ \"\\196\\160Hello\"; \"\\196\\160world\" ];\n\n (* Should NOT add extra space when text already starts with space *)\n test_case \" hello\" [ \"\\196\\160hello\" ];\n test_case \" hello\" [ \"\\196\\160\"; \"\\196\\160hello\" ]\n\nlet test_byte_level_special_chars () =\n let tokenizer = Pre.byte_level ~add_prefix_space:false ~use_regex:true () in\n\n let test_case text desc =\n let result = Pre.pre_tokenize tokenizer text in\n let pieces = List.map fst result in\n (* Just verify it doesn't crash and produces something *)\n equal\n ~msg:(Printf.sprintf \"ByteLevel handles %s\" desc)\n bool true\n (List.length pieces > 0)\n in\n\n (* Punctuation *)\n test_case \".\" \"period\";\n test_case \"!\" \"exclamation\";\n test_case \"?\" \"question\";\n test_case \",\" \"comma\";\n test_case \";\" \"semicolon\";\n test_case \":\" \"colon\";\n\n (* Special characters *)\n test_case \"@\" \"at sign\";\n test_case \"#\" \"hash\";\n test_case \"$\" \"dollar\";\n test_case \"%\" \"percent\";\n test_case \"^\" \"caret\";\n test_case \"&\" \"ampersand\";\n test_case \"*\" \"asterisk\";\n\n (* Brackets and quotes *)\n test_case \"()\" \"parentheses\";\n test_case \"[]\" \"brackets\";\n test_case \"{}\" \"braces\";\n test_case \"\\\"\\\"\" \"quotes\";\n test_case \"''\" \"single quotes\";\n\n (* Numbers *)\n test_case \"123\" \"numbers\";\n test_case \"3.14\" \"decimal\";\n test_case \"1,000\" \"number with comma\";\n\n (* Mixed *)\n test_case \"Hello, world!\" \"punctuated sentence\";\n test_case \"@user #hashtag\" \"social media\";\n test_case \"test@example.com\" \"email\";\n test_case \"https://example.com\" \"URL\";\n test_case \"function()\" \"function call\";\n test_case \"a+b=c\" \"math expression\"\n\nlet test_byte_level_unicode () =\n let tokenizer = Pre.byte_level ~add_prefix_space:false ~use_regex:true () in\n\n let test_case text desc =\n let result = Pre.pre_tokenize tokenizer text in\n let pieces = List.map fst result in\n (* Byte-level encoding should handle any Unicode by encoding bytes *)\n equal\n ~msg:(Printf.sprintf \"ByteLevel handles %s\" desc)\n bool true\n (List.length pieces > 0);\n (* Check that we can reconstruct something (even if not identical due to\n encoding) *)\n let concatenated = String.concat \"\" pieces in\n equal\n ~msg:(Printf.sprintf \"ByteLevel produces non-empty output for %s\" desc)\n bool true\n (String.length concatenated > 0)\n in\n\n (* Common accented characters *)\n test_case \"café\" \"accented e\";\n test_case \"naïve\" \"diaeresis\";\n test_case \"résumé\" \"French accents\";\n\n (* Other languages *)\n test_case \"你好\" \"Chinese\";\n test_case \"こんにちは\" \"Japanese\";\n test_case \"안녕하세요\" \"Korean\";\n test_case \"Привет\" \"Russian\";\n test_case \"مرحبا\" \"Arabic\";\n\n (* Emojis *)\n test_case \"😀\" \"emoji\";\n test_case \"👍🏻\" \"emoji with skin tone\";\n test_case \"Hello 👋 World\" \"text with emoji\"\n\nlet test_byte_level_edge_cases () =\n let tokenizer = Pre.byte_level ~add_prefix_space:false ~use_regex:true () in\n\n (* Empty string *)\n let result = Pre.pre_tokenize tokenizer \"\" in\n equal ~msg:\"Empty string\" (list string) [] (List.map fst result);\n\n (* Single character *)\n let result = Pre.pre_tokenize tokenizer \"a\" in\n check_strings \"Single char\" result [ \"a\" ];\n\n (* Only spaces - Python produces ['ĠĠĠ'] all together *)\n let result = Pre.pre_tokenize tokenizer \" \" in\n check_strings \"Only spaces\" result [ \"\\196\\160\\196\\160\\196\\160\" ];\n\n (* Only punctuation - Python keeps '...' together *)\n let result = Pre.pre_tokenize tokenizer \"...\" in\n check_strings \"Only punctuation\" result [ \"...\" ];\n\n (* Very long word *)\n let long_word = String.make 100 'a' in\n let result = Pre.pre_tokenize tokenizer long_word in\n equal ~msg:\"Long word produces single token\" int 1 (List.length result);\n\n (* Mixed whitespace *)\n let result = Pre.pre_tokenize tokenizer \"hello\\tworld\\nfoo\\rbar\" in\n equal ~msg:\"Handles tabs and newlines\" bool true (List.length result > 0)\n\nlet test_bert_pretokenizer () =\n let test_case text expected =\n let result = Pre.pre_tokenize (Pre.bert ()) text in\n check_tokenization\n (Printf.sprintf \"BERT tokenization of %S\" text)\n result expected\n in\n\n (* Basic tokenization *)\n test_case \"Hello world\" [ (\"Hello\", (0, 5)); (\"world\", (6, 11)) ];\n test_case \"Hello, world!\"\n [ (\"Hello\", (0, 5)); (\",\", (5, 6)); (\"world\", (7, 12)); (\"!\", (12, 13)) ];\n\n (* Punctuation handling *)\n test_case \"test.\" [ (\"test\", (0, 4)); (\".\", (4, 5)) ];\n test_case \"a-b\" [ (\"a\", (0, 1)); (\"-\", (1, 2)); (\"b\", (2, 3)) ];\n test_case \"it's\" [ (\"it\", (0, 2)); (\"'\", (2, 3)); (\"s\", (3, 4)) ];\n\n (* Multiple spaces *)\n test_case \"hello world\" [ (\"hello\", (0, 5)); (\"world\", (7, 12)) ];\n\n (* Unicode *)\n test_case \"café\" [ (\"café\", (0, 5)) ];\n\n (* Note: e is 2 bytes in UTF-8 *)\n\n (* Empty and whitespace *)\n test_case \"\" [];\n test_case \" \" []\n\nlet test_whitespace_pretokenizer () =\n let test_case text expected =\n let result = Pre.pre_tokenize (Pre.whitespace ()) text in\n check_tokenization\n (Printf.sprintf \"Whitespace tokenization of %S\" text)\n result expected\n in\n\n (* Pattern is \\w+|[^\\w\\s]+ *)\n test_case \"Hello world\" [ (\"Hello\", (0, 5)); (\"world\", (6, 11)) ];\n test_case \"Hello, world!\"\n [ (\"Hello\", (0, 5)); (\",\", (5, 6)); (\"world\", (7, 12)); (\"!\", (12, 13)) ];\n test_case \"test_var\" [ (\"test_var\", (0, 8)) ];\n (* underscore is part of \\w *)\n test_case \"123abc\" [ (\"123abc\", (0, 6)) ];\n (* numbers are part of \\w *)\n test_case \"a+b=c\"\n [\n (\"a\", (0, 1)); (\"+\", (1, 2)); (\"b\", (2, 3)); (\"=\", (3, 4)); (\"c\", (4, 5));\n ]\n\nlet test_whitespace_split () =\n let test_case text expected =\n let result = Pre.pre_tokenize (Pre.whitespace_split ()) text in\n check_tokenization\n (Printf.sprintf \"WhitespaceSplit of %S\" text)\n result expected\n in\n\n (* Simple split on whitespace *)\n test_case \"Hello world\" [ (\"Hello\", (0, 5)); (\"world\", (6, 11)) ];\n test_case \" Hello world \" [ (\"Hello\", (2, 7)); (\"world\", (9, 14)) ];\n test_case \"one\\ttwo\\nthree\"\n [ (\"one\", (0, 3)); (\"two\", (4, 7)); (\"three\", (8, 13)) ];\n test_case \"\" [];\n test_case \" \" []\n\nlet test_punctuation_pretokenizer () =\n (* Test different behaviors *)\n let test_isolated text expected =\n let tokenizer = Pre.punctuation ~behavior:`Isolated () in\n let result = Pre.pre_tokenize tokenizer text in\n check_tokenization\n (Printf.sprintf \"Punctuation Isolated %S\" text)\n result expected\n in\n\n let test_removed text expected =\n let tokenizer = Pre.punctuation ~behavior:`Removed () in\n let result = Pre.pre_tokenize tokenizer text in\n check_tokenization\n (Printf.sprintf \"Punctuation Removed %S\" text)\n result expected\n in\n\n (* Isolated behavior *)\n test_isolated \"Hello, world!\"\n [ (\"Hello\", (0, 5)); (\",\", (5, 6)); (\" world\", (6, 12)); (\"!\", (12, 13)) ];\n\n (* Removed behavior *)\n test_removed \"Hello, world!\" [ (\"Hello\", (0, 5)); (\" world\", (6, 12)) ];\n\n (* Multiple punctuation *)\n test_isolated \"test...end\"\n [\n (\"test\", (0, 4));\n (\".\", (4, 5));\n (\".\", (5, 6));\n (\".\", (6, 7));\n (\"end\", (7, 10));\n ];\n\n (* Unicode punctuation *)\n test_isolated \"test—end\" [ (\"test\", (0, 4)); (\"—\", (4, 7)); (\"end\", (7, 10)) ]\n(* em dash is 3 bytes *)\n\nlet test_digits_pretokenizer () =\n let test_individual text expected =\n let tokenizer = Pre.digits ~individual_digits:true () in\n let result = Pre.pre_tokenize tokenizer text in\n check_tokenization\n (Printf.sprintf \"Digits individual %S\" text)\n result expected\n in\n\n let test_grouped text expected =\n let tokenizer = Pre.digits ~individual_digits:false () in\n let result = Pre.pre_tokenize tokenizer text in\n check_tokenization (Printf.sprintf \"Digits grouped %S\" text) result expected\n in\n\n (* Individual digits *)\n test_individual \"123\" [ (\"1\", (0, 1)); (\"2\", (1, 2)); (\"3\", (2, 3)) ];\n test_individual \"a1b2\"\n [ (\"a\", (0, 1)); (\"1\", (1, 2)); (\"b\", (2, 3)); (\"2\", (3, 4)) ];\n\n (* Grouped digits *)\n test_grouped \"123\" [ (\"123\", (0, 3)) ];\n test_grouped \"a123b456\"\n [ (\"a\", (0, 1)); (\"123\", (1, 4)); (\"b\", (4, 5)); (\"456\", (5, 8)) ];\n test_grouped \"3.14\" [ (\"3\", (0, 1)); (\".\", (1, 2)); (\"14\", (2, 4)) ]\n\nlet test_split_pretokenizer () =\n let test_case pattern behavior text expected =\n let tokenizer = Pre.split ~pattern ~behavior () in\n let result = Pre.pre_tokenize tokenizer text in\n check_tokenization\n (Printf.sprintf \"Split pattern=%S behavior=%s text=%S\" pattern\n (match behavior with\n | `Isolated -> \"Isolated\"\n | `Removed -> \"Removed\"\n | `Merged_with_previous -> \"MergedPrev\"\n | `Merged_with_next -> \"MergedNext\"\n | `Contiguous -> \"Contiguous\")\n text)\n result expected\n in\n\n (* Test different behaviors *)\n test_case \",\" `Isolated \"a,b,c\"\n [\n (\"a\", (0, 1)); (\",\", (1, 2)); (\"b\", (2, 3)); (\",\", (3, 4)); (\"c\", (4, 5));\n ];\n\n test_case \",\" `Removed \"a,b,c\" [ (\"a\", (0, 1)); (\"b\", (2, 3)); (\"c\", (4, 5)) ];\n\n test_case \",\" `Merged_with_previous \"a,b,c\"\n [ (\"a,\", (0, 2)); (\"b,\", (2, 4)); (\"c\", (4, 5)) ];\n\n test_case \",\" `Merged_with_next \"a,b,c\"\n [ (\"a\", (0, 1)); (\",b\", (1, 3)); (\",c\", (3, 5)) ];\n\n (* Test with longer pattern *)\n test_case \"::\" `Isolated \"a::b::c\"\n [\n (\"a\", (0, 1)); (\"::\", (1, 3)); (\"b\", (3, 4)); (\"::\", (4, 6)); (\"c\", (6, 7));\n ]\n\nlet test_char_delimiter_split () =\n let test_case delim text expected =\n let result = Pre.pre_tokenize (Pre.char_delimiter delim) text in\n check_tokenization\n (Printf.sprintf \"CharDelimiterSplit delim='%c' text=%S\" delim text)\n result expected\n in\n\n test_case ',' \"a,b,c\" [ (\"a\", (0, 1)); (\"b\", (2, 3)); (\"c\", (4, 5)) ];\n test_case ' ' \"hello world\" [ (\"hello\", (0, 5)); (\"world\", (6, 11)) ];\n test_case '|' \"one|two|three\"\n [ (\"one\", (0, 3)); (\"two\", (4, 7)); (\"three\", (8, 13)) ];\n test_case ',' \"\" [];\n test_case ',' \",\" []\n\nlet test_sequence_pretokenizer () =\n (* Combine whitespace split then punctuation isolation *)\n let tokenizers =\n [ Pre.whitespace_split (); Pre.punctuation ~behavior:`Isolated () ]\n in\n let tokenizer = Pre.sequence tokenizers in\n\n let test_case text expected =\n let result = Pre.pre_tokenize tokenizer text in\n check_tokenization (Printf.sprintf \"Sequence %S\" text) result expected\n in\n\n (* First splits on whitespace, then isolates punctuation in each piece *)\n test_case \"Hello, world!\"\n [ (\"Hello\", (0, 5)); (\",\", (5, 6)); (\"world\", (7, 12)); (\"!\", (12, 13)) ];\n\n (* Multiple words and punctuation *)\n test_case \"test. another, example!\"\n [\n (\"test\", (0, 4));\n (\".\", (4, 5));\n (\"another\", (6, 13));\n (\",\", (13, 14));\n (\"example\", (15, 22));\n (\"!\", (22, 23));\n ]\n\nlet test_fixed_length () =\n let test_case length text expected =\n let result = Pre.pre_tokenize (Pre.fixed_length length) text in\n check_tokenization\n (Printf.sprintf \"FixedLength %d %S\" length text)\n result expected\n in\n\n test_case 3 \"abcdefghi\" [ (\"abc\", (0, 3)); (\"def\", (3, 6)); (\"ghi\", (6, 9)) ];\n test_case 2 \"abcde\" [ (\"ab\", (0, 2)); (\"cd\", (2, 4)); (\"e\", (4, 5)) ];\n test_case 5 \"hello\" [ (\"hello\", (0, 5)) ];\n test_case 0 \"test\" [];\n test_case 3 \"\" [];\n\n (* With UTF-8 - counts characters not bytes *)\n test_case 2 \"café\" [ (\"ca\", (0, 2)); (\"fé\", (2, 5)) ]\n(* e is 2 bytes *)\n\nlet test_unicode_scripts () =\n let test_case text desc =\n let tokenizer = Pre.unicode_scripts () in\n let result = Pre.pre_tokenize tokenizer text in\n (* Just verify it runs without crashing and produces something reasonable *)\n equal\n ~msg:(Printf.sprintf \"UnicodeScripts %s\" desc)\n bool true\n (List.length result >= 0)\n in\n\n test_case \"Hello world\" \"Latin text\";\n test_case \"Hello世界\" \"Mixed Latin and Chinese\";\n test_case \"Привет мир\" \"Cyrillic\";\n test_case \"مرحبا بالعالم\" \"Arabic\";\n test_case \"こんにちは世界\" \"Japanese\";\n test_case \"\" \"Empty string\"\n\nlet test_metaspace_basic () =\n let test_case text expected =\n let result =\n Pre.pre_tokenize\n (Pre.metaspace ~replacement:'_' ~prepend_scheme:`Always ~split:true ())\n text\n in\n check_strings (Printf.sprintf \"Metaspace %S\" text) result expected\n in\n\n test_case \"Hello world\" [ \"_Hello\"; \"_world\" ];\n test_case \" starts with space\" [ \"_starts\"; \"_with\"; \"_space\" ];\n test_case \"\" []\n\nlet () =\n run \"Pre-tokenizers Test Suite\"\n [\n group \"byte_level\"\n [\n test \"ByteLevel basic\" test_byte_level_basic;\n test \"ByteLevel prefix space\" test_byte_level_prefix_space;\n test \"ByteLevel special chars\" test_byte_level_special_chars;\n test \"ByteLevel unicode\" test_byte_level_unicode;\n test \"ByteLevel edge cases\" test_byte_level_edge_cases;\n ];\n group \"bert\" [ test \"BERT tokenization\" test_bert_pretokenizer ];\n group \"whitespace\"\n [\n test \"Whitespace tokenization\" test_whitespace_pretokenizer;\n test \"WhitespaceSplit\" test_whitespace_split;\n ];\n group \"punctuation\"\n [ test \"Punctuation behaviors\" test_punctuation_pretokenizer ];\n group \"digits\" [ test \"Digits tokenization\" test_digits_pretokenizer ];\n group \"split\"\n [\n test \"Split with patterns\" test_split_pretokenizer;\n test \"CharDelimiterSplit\" test_char_delimiter_split;\n ];\n group \"sequence\"\n [ test \"Sequence of tokenizers\" test_sequence_pretokenizer ];\n group \"fixed_length\" [ test \"FixedLength chunks\" test_fixed_length ];\n group \"unicode_scripts\" [ test \"UnicodeScripts\" test_unicode_scripts ];\n group \"metaspace\" [ test \"Metaspace basic\" test_metaspace_basic ];\n ]\n"
},
{
"path": "packages/brot/test/test_processors.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nopen Brot\n\nlet make_encoding ~ids ~tokens ~type_id =\n let len = Array.length ids in\n Encoding.create ~ids:(Array.copy ids) ~type_ids:(Array.make len type_id)\n ~tokens:(Array.copy tokens) ~words:(Array.make len None)\n ~offsets:(Array.make len (0, 0))\n ~special_tokens_mask:(Array.make len 0) ~attention_mask:(Array.make len 1)\n ()\n\nlet json_obj pairs =\n Jsont.Json.object' (List.map (fun (k, v) -> (Jsont.Json.name k, v)) pairs)\n\nlet test_template_multi_special () =\n let processor =\n Result.get_ok\n (Post_processor.of_json\n (json_obj\n [\n (\"type\", Jsont.Json.string \"TemplateProcessing\");\n ( \"single\",\n Jsont.Json.list\n [\n json_obj\n [\n ( \"SpecialToken\",\n json_obj\n [\n (\"id\", Jsont.Json.string \"\");\n (\"type_id\", Jsont.Json.int 2);\n ] );\n ];\n json_obj\n [\n ( \"Sequence\",\n json_obj\n [\n (\"id\", Jsont.Json.string \"A\");\n (\"type_id\", Jsont.Json.int 0);\n ] );\n ];\n ] );\n (\"pair\", Jsont.Json.null ());\n ( \"special_tokens\",\n json_obj\n [\n ( \"\",\n json_obj\n [\n (\"id\", Jsont.Json.string \"\");\n ( \"ids\",\n Jsont.Json.list\n [ Jsont.Json.int 100; Jsont.Json.int 101 ] );\n ( \"tokens\",\n Jsont.Json.list\n [\n Jsont.Json.string \"\";\n Jsont.Json.string \"\";\n ] );\n ] );\n ] );\n ]))\n in\n let base = make_encoding ~ids:[| 10 |] ~tokens:[| \"hello\" |] ~type_id:0 in\n let encoding =\n Post_processor.process processor base ~add_special_tokens:true\n in\n equal ~msg:\"ids\" (array int) [| 100; 101; 10 |] (Encoding.ids encoding);\n equal ~msg:\"tokens\" (array string)\n [| \"\"; \"\"; \"hello\" |]\n (Encoding.tokens encoding);\n equal ~msg:\"type ids\" (array int) [| 2; 2; 0 |] (Encoding.type_ids encoding);\n equal ~msg:\"special mask\" (array int) [| 1; 1; 0 |]\n (Encoding.special_tokens_mask encoding);\n equal ~msg:\"attention mask\" (array int) [| 1; 1; 1 |]\n (Encoding.attention_mask encoding);\n equal ~msg:\"added tokens single\" int 2\n (Post_processor.added_tokens processor ~is_pair:false)\n\nlet test_template_pair_type_ids () =\n let processor =\n Post_processor.template ~single:\"$A [SEP]\"\n ~pair:\"[CLS]:0 $A:0 [SEP]:0 $B:3 [SEP]:3\"\n ~special_tokens:[ (\"[CLS]\", 101); (\"[SEP]\", 102) ]\n ()\n in\n let seq_a =\n make_encoding ~ids:[| 10; 11 |] ~tokens:[| \"hello\"; \"world\" |] ~type_id:0\n in\n let seq_b = make_encoding ~ids:[| 20 |] ~tokens:[| \"pair\" |] ~type_id:1 in\n let encoding =\n Post_processor.process processor ~pair:seq_b seq_a ~add_special_tokens:true\n in\n equal ~msg:\"pair ids\" (array int)\n [| 101; 10; 11; 102; 20; 102 |]\n (Encoding.ids encoding);\n equal ~msg:\"pair tokens\" (array string)\n [| \"[CLS]\"; \"hello\"; \"world\"; \"[SEP]\"; \"pair\"; \"[SEP]\" |]\n (Encoding.tokens encoding);\n equal ~msg:\"pair type ids\" (array int) [| 0; 0; 0; 0; 3; 3 |]\n (Encoding.type_ids encoding);\n equal ~msg:\"pair special mask\" (array int) [| 1; 0; 0; 1; 0; 1 |]\n (Encoding.special_tokens_mask encoding);\n equal ~msg:\"added tokens pair\" int 3\n (Post_processor.added_tokens processor ~is_pair:true);\n let no_special =\n Post_processor.process processor ~pair:seq_b seq_a ~add_special_tokens:false\n in\n equal ~msg:\"no-special ids\" (array int) (Encoding.ids seq_a)\n (Encoding.ids no_special)\n\nlet () =\n run \"Processors\"\n [\n group \"template\"\n [\n test \"multi-id special expansion\" test_template_multi_special;\n test \"pair template semantics\" test_template_pair_type_ids;\n ];\n ]\n"
},
{
"path": "packages/brot/test/test_tokenization.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Tokenization tests for brot *)\n\nopen Windtrap\nopen Brot\n\n(* Helper function to tokenize text *)\nlet tokenize_text text =\n (* Pre-tokenize to get all unique tokens *)\n let pre_tokens =\n Pre_tokenizer.pre_tokenize (Pre_tokenizer.whitespace ()) text\n in\n let unique_tokens =\n List.fold_left\n (fun acc (tok, _) -> if List.mem tok acc then acc else tok :: acc)\n [] pre_tokens\n |> List.rev\n in\n (* Build vocabulary with all tokens from the text plus extras *)\n let all_tokens =\n unique_tokens\n @\n (* Add numbered words for long text test *)\n List.init 1000 (fun i -> Printf.sprintf \"word%d\" i)\n in\n let vocab = List.mapi (fun i token -> (token, i)) all_tokens in\n\n (* Create WordLevel tokenizer with the vocabulary *)\n let tokenizer =\n word_level ~vocab ~unk_token:\"\" ~pre:(Pre_tokenizer.whitespace ()) ()\n in\n encode tokenizer text |> Encoding.tokens |> Array.to_list\n\n(* Basic Tokenization Tests *)\n\nlet test_tokenize_words_simple () =\n let tokens = tokenize_text \"Hello world!\" in\n equal ~msg:\"simple words\" (list string) [ \"Hello\"; \"world\"; \"!\" ] tokens\n\nlet test_tokenize_words_punctuation () =\n let tokens = tokenize_text \"don't stop, it's fun!\" in\n equal ~msg:\"words with punctuation\" (list string)\n [ \"don\"; \"'\"; \"t\"; \"stop\"; \",\"; \"it\"; \"'\"; \"s\"; \"fun\"; \"!\" ]\n tokens\n\nlet test_tokenize_words_numbers () =\n let tokens = tokenize_text \"I have 42 apples and 3.14 pies\" in\n equal ~msg:\"words with numbers\" (list string)\n [ \"I\"; \"have\"; \"42\"; \"apples\"; \"and\"; \"3\"; \".\"; \"14\"; \"pies\" ]\n tokens\n\nlet test_tokenize_words_empty () =\n let tokens = tokenize_text \"\" in\n equal ~msg:\"empty string\" (list string) [] tokens\n\nlet test_tokenize_words_whitespace_only () =\n let tokens = tokenize_text \" \\t\\n \" in\n equal ~msg:\"whitespace only\" (list string) [] tokens\n\nlet test_tokenize_words_special_chars () =\n let tokens = tokenize_text \"hello@world.com #ml $100 C++\" in\n equal ~msg:\"special characters\" (list string)\n [ \"hello\"; \"@\"; \"world\"; \".\"; \"com\"; \"#\"; \"ml\"; \"$\"; \"100\"; \"C\"; \"++\" ]\n tokens\n\n(* Character Tokenization Tests *)\n\nlet tokenize_chars text =\n let chars = ref [] in\n String.iter (fun c -> chars := String.make 1 c :: !chars) text;\n List.rev !chars\n\nlet test_tokenize_chars_ascii () =\n let tokens = tokenize_chars \"Hi!\" in\n equal ~msg:\"ASCII chars\" (list string) [ \"H\"; \"i\"; \"!\" ] tokens\n\nlet test_tokenize_chars_unicode () =\n let tokens = tokenize_chars \"Hello 👋 世界\" in\n (* Note: UTF-8 encoding means multi-byte chars may appear differently *)\n equal ~msg:\"has tokens\" bool true (List.length tokens > 0)\n\nlet test_tokenize_chars_empty () =\n let tokens = tokenize_chars \"\" in\n equal ~msg:\"empty string chars\" (list string) [] tokens\n\n(* Pre-tokenizer Pattern Tests *)\n\nlet test_tokenize_regex_words () =\n (* Use the helper that sets up vocabulary properly *)\n let tokens = tokenize_text \"hello-world test_123\" in\n equal ~msg:\"regex words\" (list string)\n [ \"hello\"; \"-\"; \"world\"; \"test_123\" ]\n tokens\n\nlet test_tokenize_regex_custom () =\n (* Test with punctuation pre-tokenizer *)\n let text = \"don't stop!\" in\n let pre_tokens =\n Pre_tokenizer.pre_tokenize (Pre_tokenizer.punctuation ()) text\n in\n let vocab = List.mapi (fun i (tok, _) -> (tok, i)) pre_tokens in\n let tokenizer =\n word_level ~vocab ~unk_token:\"\" ~pre:(Pre_tokenizer.punctuation ()) ()\n in\n let tokens = encode tokenizer text |> Encoding.tokens |> Array.to_list in\n equal ~msg:\"has tokens\" bool true (List.length tokens > 0)\n\nlet test_tokenize_regex_no_match () =\n let tokenizer = word_level () in\n let tokens =\n encode tokenizer \"no numbers here\" |> Encoding.tokens |> Array.to_list\n in\n equal ~msg:\"regex no match\" (list string) [] tokens\n\n(* Unigram Model Tests *)\n\n(* Round-trip lookups *)\nlet test_unigram_roundtrip () =\n let tokens = [ \"hello\"; \"world\"; \"test\" ] in\n let vocab = List.map (fun token -> (token, 0.0)) tokens in\n let tokenizer = unigram ~vocab () in\n List.iteri\n (fun expected_id token ->\n equal\n ~msg:(Printf.sprintf \"token_to_id '%s'\" token)\n (option int) (Some expected_id)\n (token_to_id tokenizer token);\n equal\n ~msg:(Printf.sprintf \"id_to_token %d\" expected_id)\n (option string) (Some token)\n (id_to_token tokenizer expected_id))\n tokens\n\n(* token_to_id - out of vocab *)\nlet test_unigram_token_to_id_oov () =\n let tokenizer = unigram ~vocab:[ (\"hello\", 0.0); (\"world\", 0.0) ] () in\n equal ~msg:\"token_to_id out-of-vocab\" (option int) None\n (token_to_id tokenizer \"missing\")\n\n(* id_to_token - out of bounds *)\nlet test_unigram_id_to_token_oob () =\n let tokenizer = unigram ~vocab:[ (\"hello\", 0.0); (\"world\", 0.0) ] () in\n equal ~msg:\"id_to_token negative\" (option string) None\n (id_to_token tokenizer (-1));\n equal ~msg:\"id_to_token out of bounds\" (option string) None\n (id_to_token tokenizer 10)\n\n(* Test empty vocabulary *)\nlet test_unigram_empty_vocab () =\n let tokenizer = unigram ~vocab:[] () in\n equal ~msg:\"empty vocab token_to_id\" (option int) None\n (token_to_id tokenizer \"test\");\n equal ~msg:\"empty vocab id_to_token\" (option string) None\n (id_to_token tokenizer 0)\n\n(* Test special characters and unicode *)\nlet test_unigram_special_tokens () =\n let tokenizer =\n unigram\n ~vocab:\n [\n (\"\", 0.0);\n (\"\", 0.0);\n (\" \", 0.0);\n (\"▁hello\", 0.0);\n (\"世界\", 0.0);\n ]\n ()\n in\n equal ~msg:\"special \" (option int) (Some 0)\n (token_to_id tokenizer \"\");\n equal ~msg:\"special \" (option int) (Some 1) (token_to_id tokenizer \"\");\n equal ~msg:\"sentencepiece token\" (option int) (Some 3)\n (token_to_id tokenizer \"▁hello\");\n equal ~msg:\"unicode token\" (option int) (Some 4) (token_to_id tokenizer \"世界\");\n equal ~msg:\"id to unicode\" (option string) (Some \"世界\")\n (id_to_token tokenizer 4)\n\nlet test_unigram_encode_sequence () =\n let tokenizer = unigram ~vocab:[ (\"hello\", 0.0); (\"world\", 0.0) ] () in\n let encoding = encode tokenizer \"hello world\" in\n let tokens = Encoding.tokens encoding |> Array.to_list in\n equal ~msg:\"unigram encode tokens\" (list string) [ \"hello\"; \"world\" ] tokens\n\nlet test_pad_token_set_at_construction () =\n let vocab = [ (\"hello\", 0); (\"world\", 1); (\"\", 2); (\"[PAD]\", 3) ] in\n let tokenizer =\n word_level ~vocab ~unk_token:\"\"\n ~pre:(Pre_tokenizer.whitespace ())\n ~specials:[ special \"[PAD]\" ]\n ~pad_token:\"[PAD]\" ()\n in\n equal ~msg:\"pad token set\" (option string) (Some \"[PAD]\")\n (pad_token tokenizer);\n let pad_id =\n match token_to_id tokenizer \"[PAD]\" with\n | Some id -> id\n | None -> failwith \"missing pad id\"\n in\n let encoding =\n encode tokenizer \"hello\"\n ~padding:\n {\n length = `Fixed 3;\n direction = `Right;\n pad_id = None;\n pad_type_id = None;\n pad_token = None;\n }\n in\n let ids = Encoding.ids encoding |> Array.to_list in\n let pad_ids = List.tl ids in\n equal ~msg:\"pad id matches configured token\" (list int) [ pad_id; pad_id ]\n pad_ids\n\n(* Edge Cases *)\n\nlet test_tokenize_long_text () =\n let text =\n String.concat \" \" (List.init 1000 (fun i -> Printf.sprintf \"word%d\" i))\n in\n let tokens = tokenize_text text in\n equal ~msg:\"long text token count\" int 1000 (List.length tokens)\n\nlet test_tokenize_repeated_punctuation () =\n let tokens = tokenize_text \"wow!!! really???\" in\n equal ~msg:\"repeated punctuation\" (list string)\n [ \"wow\"; \"!!!\"; \"really\"; \"???\" ]\n tokens\n\nlet test_tokenize_mixed_whitespace () =\n let tokens = tokenize_text \"hello\\tworld\\nthere\\r\\nfriend\" in\n equal ~msg:\"mixed whitespace\" (list string)\n [ \"hello\"; \"world\"; \"there\"; \"friend\" ]\n tokens\n\n(* Test Suite *)\n\nlet tokenization_tests =\n [\n (* Words tokenization *)\n test \"tokenize words simple\" test_tokenize_words_simple;\n test \"tokenize words punctuation\" test_tokenize_words_punctuation;\n test \"tokenize words numbers\" test_tokenize_words_numbers;\n test \"tokenize words empty\" test_tokenize_words_empty;\n test \"tokenize words whitespace only\" test_tokenize_words_whitespace_only;\n test \"tokenize words special chars\" test_tokenize_words_special_chars;\n (* Character tokenization *)\n test \"tokenize chars ASCII\" test_tokenize_chars_ascii;\n test \"tokenize chars unicode\" test_tokenize_chars_unicode;\n test \"tokenize chars empty\" test_tokenize_chars_empty;\n (* Regex tokenization *)\n test \"tokenize regex words\" test_tokenize_regex_words;\n test \"tokenize regex custom\" test_tokenize_regex_custom;\n test \"tokenize regex no match\" test_tokenize_regex_no_match;\n (* Edge cases *)\n test \"tokenize long text\" test_tokenize_long_text;\n test \"tokenize repeated punctuation\" test_tokenize_repeated_punctuation;\n test \"tokenize mixed whitespace\" test_tokenize_mixed_whitespace;\n (* Unigram model tests *)\n test \"unigram roundtrip\" test_unigram_roundtrip;\n test \"unigram token_to_id out-of-vocab\" test_unigram_token_to_id_oov;\n test \"unigram id_to_token out-of-bounds\" test_unigram_id_to_token_oob;\n test \"unigram empty vocab\" test_unigram_empty_vocab;\n test \"unigram special tokens\" test_unigram_special_tokens;\n test \"unigram encode sequence\" test_unigram_encode_sequence;\n test \"pad token reassignment updates id\" test_pad_token_set_at_construction;\n ]\n\nlet () = run \"brot tokenization\" [ group \"tokenization\" tokenization_tests ]\n"
},
{
"path": "packages/brot/test/test_unicode.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Unicode processing tests for brot *)\n\nopen Windtrap\nopen Brot\n\n(* Normalization via public API *)\n\nlet test_lowercase_normalization () =\n let text = \"HELLO WORLD\" in\n let normalizer = Normalizer.lowercase in\n let result = Normalizer.apply normalizer text in\n equal ~msg:\"lowercase\" string \"hello world\" result\n\nlet test_strip_accents_normalization () =\n let text = \"caf\\xC3\\xA9 na\\xC3\\xAFve r\\xC3\\xA9sum\\xC3\\xA9\" in\n let normalizer =\n Normalizer.sequence [ Normalizer.nfd; Normalizer.strip_accents ]\n in\n let result = Normalizer.apply normalizer text in\n equal ~msg:\"strip accents\" string \"cafe naive resume\" result\n\nlet test_normalization_sequence () =\n let text = \" HELLO World \" in\n let normalizer =\n Normalizer.sequence\n [\n Normalizer.lowercase;\n Normalizer.strip ();\n Normalizer.replace ~pattern:\"\\\\s+\" ~replacement:\" \";\n ]\n in\n let result = Normalizer.apply normalizer text in\n equal ~msg:\"sequence\" string \"hello world\" result\n\n(* Integration with Tokenizer *)\n\nlet test_tokenize_with_normalization () =\n let text = \"HELLO WORLD!\" in\n let normalizer =\n Normalizer.sequence\n [\n Normalizer.lowercase;\n Normalizer.replace ~pattern:\"\\\\s+\" ~replacement:\" \";\n ]\n in\n let tokenizer =\n word_level ~normalizer ~pre:(Pre_tokenizer.whitespace ()) ()\n in\n let tokenizer = add_tokens tokenizer [ \"hello\"; \"world\"; \"!\" ] in\n let tokens = encode tokenizer text |> Encoding.tokens |> Array.to_list in\n equal ~msg:\"normalized tokenization\" (list string) [ \"hello\"; \"world\"; \"!\" ]\n tokens\n\nlet test_tokenize_unicode_words () =\n let text = \"café résumé naïve\" in\n let tokenizer = word_level ~pre:(Pre_tokenizer.whitespace ()) () in\n let tokenizer = add_tokens tokenizer [ \"café\"; \"résumé\"; \"naïve\" ] in\n let tokens = encode tokenizer text |> Encoding.tokens |> Array.to_list in\n equal ~msg:\"tokenized unicode\" bool true (List.length tokens > 0)\n\nlet test_malformed_unicode () =\n let text = \"Hello\" ^ String.make 1 '\\xFF' ^ String.make 1 '\\xFE' ^ \"World\" in\n let tokenizer = chars () in\n let tokens = encode tokenizer text |> Encoding.tokens |> Array.to_list in\n equal ~msg:\"handled malformed\" bool true (List.length tokens > 0)\n\n(* Test Suite *)\n\nlet unicode_tests =\n [\n (* Normalization *)\n test \"lowercase normalization\" test_lowercase_normalization;\n test \"strip accents normalization\" test_strip_accents_normalization;\n test \"normalization sequence\" test_normalization_sequence;\n (* Integration *)\n test \"tokenize with normalization\" test_tokenize_with_normalization;\n test \"tokenize unicode words\" test_tokenize_unicode_words;\n (* Error handling *)\n test \"malformed unicode\" test_malformed_unicode;\n ]\n\nlet () = run \"brot unicode\" [ group \"unicode\" unicode_tests ]\n"
},
{
"path": "packages/brot/test/test_vocab.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nopen Brot\n\nlet test_vocab_create_empty () =\n let tokenizer = word_level () in\n let vocab = vocab tokenizer in\n equal ~msg:\"empty vocab size\" int 0 (List.length vocab)\n\nlet test_vocab_with_tokenizer () =\n let tokenizer = word_level () in\n let vocab = vocab tokenizer in\n equal ~msg:\"initial vocab size\" int 0 (List.length vocab)\n\nlet test_vocab_add_tokens () =\n let tokenizer =\n add_tokens\n (word_level ~specials:[ special \"\"; special \"\" ] ())\n [ \"hello\"; \"world\" ]\n in\n let vocab_size = vocab_size tokenizer in\n equal ~msg:\"vocab size increased\" bool true (vocab_size >= 2)\n\nlet test_vocab_encode_decode () =\n let tokenizer =\n add_tokens\n (word_level ~pre:(Pre_tokenizer.whitespace ()) ())\n [ \"hello\"; \"world\" ]\n in\n let ids = encode tokenizer \"hello world\" |> Encoding.ids in\n equal ~msg:\"encoded ids\" bool true (Array.length ids > 0);\n let decoded = decode tokenizer ids in\n equal ~msg:\"decoded text\" string \"hello world\" decoded\n\nlet test_vocab_batch_encode () =\n let tokenizer = add_tokens (Brot.word_level ()) [ \"hello\"; \"world\" ] in\n let encodings = encode_batch tokenizer [ \"hello\"; \"world\" ] in\n equal ~msg:\"batch size\" int 2 (List.length encodings)\n\nlet test_vocab_special_tokens () =\n let tokenizer =\n add_tokens\n (word_level ~specials:[ special \"[CLS]\"; special \"[SEP]\" ] ())\n [ \"test\" ]\n in\n let tokens =\n encode ~add_special_tokens:true tokenizer \"test\" |> Encoding.tokens\n in\n equal ~msg:\"tokens emitted\" bool true (Array.length tokens > 0)\n\nlet test_vocab_save_load () =\n let tokenizer =\n add_tokens (Brot.word_level ()) [ \"hello\"; \"world\"; \"test\" ]\n in\n let json = to_json tokenizer in\n match from_json json with\n | Error msg -> failf \"failed to round-trip tokenizer: %s\" msg\n | Ok reloaded ->\n let original_vocab = vocab tokenizer in\n let loaded_vocab = vocab reloaded in\n equal ~msg:\"vocab size matches\" int\n (List.length original_vocab)\n (List.length loaded_vocab);\n List.iter\n (fun (token, _) ->\n equal\n ~msg:(Printf.sprintf \"token %s preserved\" token)\n bool true\n (Option.is_some (token_to_id reloaded token)))\n original_vocab\n\nlet suite =\n [\n test \"create empty\" test_vocab_create_empty;\n test \"with tokenizer\" test_vocab_with_tokenizer;\n test \"add tokens\" test_vocab_add_tokens;\n test \"encode decode\" test_vocab_encode_decode;\n test \"batch encode\" test_vocab_batch_encode;\n test \"special tokens\" test_vocab_special_tokens;\n test \"save load\" test_vocab_save_load;\n ]\n\nlet () = run \"Vocabulary tests\" [ group \"vocab\" suite ]\n"
},
{
"path": "packages/brot/test/test_wordpiece.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nopen Brot\n\nlet test_wordpiece_basic () =\n (* Create a simple vocabulary *)\n let vocab =\n [\n (\"[UNK]\", 0);\n (\"hello\", 1);\n (\"world\", 2);\n (\"##llo\", 3);\n (\"##rld\", 4);\n (\"he\", 5);\n (\"wo\", 6);\n ]\n in\n\n let tokenizer =\n wordpiece ~vocab ~unk_token:\"[UNK]\" ~continuing_subword_prefix:\"##\" ()\n in\n\n (* Test tokenizing a known word *)\n let encoding = encode tokenizer \"hello\" in\n let tokens = Encoding.tokens encoding in\n equal ~msg:\"single token for 'hello'\" int 1 (Array.length tokens);\n equal ~msg:\"token value\" string \"hello\" tokens.(0);\n\n Printf.printf \"Tokenized 'hello': \";\n Array.iter (Printf.printf \"%s \") tokens;\n Printf.printf \"\\n\";\n\n equal ~msg:\"vocabulary size\" int 7 (vocab_size tokenizer)\n\nlet test_wordpiece_subwords () =\n (* Create vocabulary with subword pieces *)\n let vocab =\n [\n (\"[UNK]\", 0);\n (\"un\", 1);\n (\"##able\", 2);\n (\"##happy\", 3);\n (\"play\", 4);\n (\"##ing\", 5);\n (\"##ed\", 6);\n ]\n in\n\n let tokenizer = wordpiece ~vocab ~unk_token:\"[UNK]\" () in\n\n (* Test word that can be split into subwords *)\n let encoding = encode tokenizer \"playing\" in\n let tokens = Encoding.tokens encoding in\n Printf.printf \"Tokenized 'playing': \";\n Array.iter (Printf.printf \"%s \") tokens;\n Printf.printf \"\\n\";\n\n equal ~msg:\"should split into subwords\" int 2 (Array.length tokens);\n equal ~msg:\"first token\" string \"play\" tokens.(0);\n equal ~msg:\"second token\" string \"##ing\" tokens.(1)\n\nlet test_wordpiece_unknown () =\n (* Create minimal vocabulary *)\n let vocab = [ (\"[UNK]\", 0); (\"hello\", 1) ] in\n\n let tokenizer = wordpiece ~vocab ~unk_token:\"[UNK]\" () in\n\n (* Test unknown word *)\n let encoding = encode tokenizer \"goodbye\" in\n let tokens = Encoding.tokens encoding in\n equal ~msg:\"unknown word becomes single token\" int 1 (Array.length tokens);\n equal ~msg:\"unknown token\" string \"[UNK]\" tokens.(0)\n\nlet test_wordpiece_max_chars () =\n (* Create vocabulary *)\n let vocab = [ (\"[UNK]\", 0); (\"test\", 1) ] in\n\n let tokenizer =\n wordpiece ~vocab ~unk_token:\"[UNK]\" ~max_input_chars_per_word:5 ()\n in\n\n (* Test word exceeding max chars *)\n let long_word = String.make 10 'a' in\n let encoding = encode tokenizer long_word in\n let tokens = Encoding.tokens encoding in\n equal ~msg:\"long word becomes unknown\" int 1 (Array.length tokens);\n equal ~msg:\"unknown token\" string \"[UNK]\" tokens.(0)\n\nlet test_wordpiece_save_load () =\n (* Create vocabulary *)\n let vocab =\n [\n (\"[PAD]\", 0);\n (\"[UNK]\", 1);\n (\"[CLS]\", 2);\n (\"[SEP]\", 3);\n (\"hello\", 4);\n (\"world\", 5);\n ]\n in\n\n let tokenizer = wordpiece ~vocab ~unk_token:\"[UNK]\" () in\n\n (* Save the model *)\n let temp_dir = Filename.temp_dir \"wordpiece_test\" \"\" in\n let files = save_model_files tokenizer ~folder:temp_dir () in\n\n (* Load the model *)\n let vocab_file = List.find (fun f -> Filename.check_suffix f \".txt\") files in\n let loaded_tokenizer = from_model_file ~vocab:vocab_file () in\n\n (* Test that loaded tokenizer works the same *)\n let original_tokens = encode tokenizer \"hello\" |> Encoding.tokens in\n let loaded_tokens = encode loaded_tokenizer \"hello\" |> Encoding.tokens in\n\n equal ~msg:\"same number of tokens\" int\n (Array.length original_tokens)\n (Array.length loaded_tokens);\n\n (* Clean up *)\n List.iter Sys.remove files;\n Unix.rmdir temp_dir\n\nlet test_tokenizer_integration () =\n (* Create a WordPiece tokenizer using the high-level API *)\n let vocab =\n [\n (\"[PAD]\", 0);\n (\"[UNK]\", 1);\n (\"[CLS]\", 2);\n (\"[SEP]\", 3);\n (\"hello\", 4);\n (\"world\", 5);\n (\"##ing\", 6);\n ]\n in\n let tokenizer = wordpiece ~vocab ~unk_token:\"[UNK]\" () in\n\n (* Test encoding *)\n let tokens = encode tokenizer \"hello\" |> Encoding.tokens |> Array.to_list in\n\n Printf.printf \"wordpiece result: \";\n List.iter (Printf.printf \"%s \") tokens;\n Printf.printf \"\\n\";\n\n equal ~msg:\"tokenizer produces output\" bool true (List.length tokens > 0)\n\nlet test_wordpiece_greedy_matching () =\n (* Test the greedy longest-match-first algorithm *)\n let vocab =\n [\n (\"[UNK]\", 0);\n (\"un\", 1);\n (\"able\", 2);\n (\"unable\", 3);\n (* Longer match should be preferred *)\n (\"##able\", 4);\n ]\n in\n\n let tokenizer = wordpiece ~vocab ~unk_token:\"[UNK]\" () in\n\n (* Should match \"unable\" as a single token, not \"un\" + \"##able\" *)\n let encoding = encode tokenizer \"unable\" in\n let tokens = Encoding.tokens encoding in\n equal ~msg:\"greedy match finds longest token\" int 1 (Array.length tokens);\n equal ~msg:\"matched full word\" string \"unable\" tokens.(0)\n\nlet () =\n run \"WordPiece tests\"\n [\n group \"basic\"\n [\n test \"basic tokenization\" test_wordpiece_basic;\n test \"subword tokenization\" test_wordpiece_subwords;\n test \"unknown tokens\" test_wordpiece_unknown;\n test \"max input chars\" test_wordpiece_max_chars;\n test \"save and load\" test_wordpiece_save_load;\n test \"tokenizer integration\" test_tokenizer_integration;\n test \"greedy matching\" test_wordpiece_greedy_matching;\n ];\n ]\n"
},
{
"path": "packages/dune",
"content": "(dirs :standard \\ nx-oxcaml)\n"
},
{
"path": "packages/fehu/README.md",
"content": "# Fehu\n\nReinforcement learning environment toolkit for OCaml, built on [Rune](../rune/)\n\nFehu provides type-safe environments, composable wrappers, trajectory\ncollection, replay buffers, GAE computation, policy evaluation, and\nvectorized environments. It follows the Gymnasium interface pattern:\nenvironments expose `reset` and `step` with typed observation and action\nspaces.\n\n## Quick Start\n\nCreate an environment, run a random policy, and evaluate:\n\n```ocaml\nopen Fehu\n\nlet () =\n let rng = Rune.Rng.key 42 in\n let env = Fehu_envs.Cartpole.make ~rng () in\n\n (* Run one episode *)\n let _obs, _info = Env.reset env () in\n let done_ = ref false in\n let total_reward = ref 0.0 in\n while not !done_ do\n let act, _ = Space.sample (Env.action_space env)\n ~rng:(Env.take_rng env) in\n let s = Env.step env act in\n total_reward := !total_reward +. s.reward;\n done_ := s.terminated || s.truncated\n done;\n Printf.printf \"Episode reward: %.0f\\n\" !total_reward;\n\n (* Evaluate over 10 episodes *)\n let stats = Eval.run env\n ~policy:(fun _obs ->\n let act, _ = Space.sample (Env.action_space env)\n ~rng:(Env.take_rng env) in act)\n ~n_episodes:10 ()\n in\n Printf.printf \"Mean reward: %.1f (std: %.1f)\\n\"\n stats.mean_reward stats.std_reward\n```\n\n## Features\n\n- **Environments**: typed `('obs, 'act, 'render) Env.t` with lifecycle enforcement (reset before step, auto-guard on terminal states)\n- **Spaces**: Discrete, Box, Multi_binary, Multi_discrete, Tuple, Dict, Sequence, Text with sampling, validation, and serialization\n- **Wrappers**: `map_observation`, `map_action`, `map_reward`, `clip_action`, `clip_observation`, `time_limit`, and custom wrappers via `Env.wrap`\n- **Trajectory collection**: `Collect.rollout` and `Collect.episodes` in structure-of-arrays form with automatic episode resets\n- **Replay buffers**: fixed-capacity circular buffer with uniform random sampling (`Buffer.sample`, `Buffer.sample_arrays`)\n- **GAE**: generalized advantage estimation with proper terminated/truncated handling (`Gae.compute`, `Gae.returns`)\n- **Evaluation**: `Eval.run` computes mean/std reward and episode length over multiple episodes\n- **Vectorized environments**: `Vec_env.create` runs multiple environments with batched step and auto-reset\n- **Rendering**: `Render.image` and `Render.rollout` for frame capture, `Env.on_render` for recording\n- **Built-in environments**: CartPole-v1, MountainCar-v0, GridWorld, RandomWalk\n\n## Libraries\n\n| Library | opam package | Description |\n|---------|-------------|-------------|\n| `fehu` | `fehu` | Core: environments, spaces, wrappers, collection, buffers, GAE, evaluation |\n| `fehu-envs` | `fehu.envs` | Built-in environments (CartPole, MountainCar, GridWorld, RandomWalk) |\n\n## Built-in Environments\n\n| Environment | Observation | Actions | Reward | Termination |\n|-------------|------------|---------|--------|-------------|\n| CartPole | Box [4] (x, v, θ, ω) | Discrete 2 | +1.0 per step | Pole > ±12° or cart > ±2.4, truncated at 500 |\n| MountainCar | Box [2] (position, velocity) | Discrete 3 | −1.0 per step | Position ≥ 0.5 with v ≥ 0, truncated at 200 |\n| GridWorld | Multi_discrete [5; 5] | Discrete 4 | +10 at goal, −1 otherwise | Reach (4,4), truncated at 200 |\n| RandomWalk | Box [1] | Discrete 2 | −|position| | None, truncated at 200 |\n\n## Contributing\n\nSee the [Raven monorepo README](../README.md) for guidelines.\n\n## License\n\nISC License. See [LICENSE](../LICENSE) for details.\n"
},
{
"path": "packages/fehu/bench/bench_fehu.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Gae = Fehu.Gae\nmodule Buffer = Fehu.Buffer\n\nlet gamma = 0.99\nlet lambda = 0.95\n\nlet make_arrays n =\n let rewards = Array.init n (fun i -> Float.of_int (i mod 5)) in\n let values = Array.init n (fun i -> Float.of_int (i mod 10) *. 0.1) in\n let terminated = Array.init n (fun i -> i mod 50 = 49) in\n let truncated = Array.init n (fun i -> i mod 100 = 99) in\n let next_values =\n Array.init n (fun i -> Float.of_int ((i + 1) mod 10) *. 0.1)\n in\n (rewards, values, terminated, truncated, next_values)\n\nlet gae_benchmarks () =\n let sizes = [ (\"256\", 256); (\"1024\", 1024); (\"4096\", 4096) ] in\n let benches = ref [] in\n List.iter\n (fun (label, n) ->\n let rewards, values, terminated, truncated, next_values = make_arrays n in\n benches :=\n Thumper.bench (Printf.sprintf \"compute n=%s\" label) (fun () ->\n Gae.compute ~rewards ~values ~terminated ~truncated ~next_values\n ~gamma ~lambda)\n :: !benches)\n sizes;\n List.rev !benches\n\nlet gae_from_values_benchmarks () =\n let sizes = [ (\"256\", 256); (\"1024\", 1024); (\"4096\", 4096) ] in\n let benches = ref [] in\n List.iter\n (fun (label, n) ->\n let rewards, values, terminated, truncated, _ = make_arrays n in\n benches :=\n Thumper.bench (Printf.sprintf \"compute_from_values n=%s\" label)\n (fun () ->\n Gae.compute_from_values ~rewards ~values ~terminated ~truncated\n ~last_value:0.0 ~gamma ~lambda)\n :: !benches)\n sizes;\n List.rev !benches\n\nlet returns_benchmarks () =\n let sizes = [ (\"256\", 256); (\"1024\", 1024); (\"4096\", 4096) ] in\n let benches = ref [] in\n List.iter\n (fun (label, n) ->\n let rewards, _, terminated, truncated, _ = make_arrays n in\n benches :=\n Thumper.bench (Printf.sprintf \"returns n=%s\" label) (fun () ->\n Gae.returns ~rewards ~terminated ~truncated ~gamma)\n :: !benches)\n sizes;\n List.rev !benches\n\nlet normalize_benchmarks () =\n let sizes = [ (\"256\", 256); (\"1024\", 1024); (\"4096\", 4096) ] in\n let benches = ref [] in\n List.iter\n (fun (label, n) ->\n let arr = Array.init n (fun i -> Float.of_int i *. 0.01) in\n benches :=\n Thumper.bench (Printf.sprintf \"normalize n=%s\" label) (fun () ->\n Gae.normalize arr)\n :: !benches)\n sizes;\n List.rev !benches\n\nlet fill_buffer capacity =\n let buf : (float, float) Buffer.t = Buffer.create ~capacity in\n for i = 0 to capacity - 1 do\n Buffer.add buf\n {\n Buffer.observation = Float.of_int i;\n action = Float.of_int (i mod 4);\n reward = Float.of_int (i mod 10) *. 0.1;\n next_observation = Float.of_int (i + 1);\n terminated = i mod 50 = 49;\n truncated = i mod 100 = 99;\n }\n done;\n buf\n\nlet buffer_add_benchmarks () =\n let capacity = 10000 in\n let buf = fill_buffer capacity in\n let tr =\n {\n Buffer.observation = 0.0;\n action = 0.0;\n reward = 1.0;\n next_observation = 1.0;\n terminated = false;\n truncated = false;\n }\n in\n [ Thumper.bench \"add (full buffer, cap=10000)\" (fun () -> Buffer.add buf tr) ]\n\nlet buffer_create_benchmarks () =\n let sizes = [ (\"100\", 100); (\"1000\", 1000); (\"10000\", 10000) ] in\n List.map\n (fun (label, n) ->\n Thumper.bench (Printf.sprintf \"create+fill cap=%s\" label) (fun () ->\n fill_buffer n))\n sizes\n\nlet build_benchmarks () =\n [\n Thumper.group \"GAE\" (gae_benchmarks ());\n Thumper.group \"GAE from values\" (gae_from_values_benchmarks ());\n Thumper.group \"Returns\" (returns_benchmarks ());\n Thumper.group \"Normalize\" (normalize_benchmarks ());\n Thumper.group \"Buffer add\" (buffer_add_benchmarks ());\n Thumper.group \"Buffer create\" (buffer_create_benchmarks ());\n ]\n\nlet () =\n let benchmarks = build_benchmarks () in\n Thumper.run \"fehu\" benchmarks\n"
},
{
"path": "packages/fehu/bench/dune",
"content": "(executable\n (name bench_fehu)\n (libraries nx fehu thumper))\n\n(rule\n (alias runtest)\n (action\n (progn\n (run %{exe:bench_fehu.exe} -q)\n (diff? fehu.thumper fehu.thumper.corrected))))\n"
},
{
"path": "packages/fehu/bench/fehu.thumper",
"content": "# thumper baseline\n# version: 1\n# suite_name: fehu\n# host: 1480401c3b76ed18\n# cpu: Apple M1 Max\n# ocaml: 5.4.1\n# git: 31747323\n# dirty: true\n# command: /Users/tmattio/Workspace/raven/_build/default/packages/fehu/bench/bench_fehu.exe --bless --quick\n\nbuffer_add/add__full_buffer__cap_10000_\talloc_words\t0.000000e+00\t0.000000e+00\t0.000000e+00\tinf\t5\t0\nbuffer_add/add__full_buffer__cap_10000_\tcpu_time\t1.226649e-08\t1.223096e-08\t1.229882e-08\t2.766164e-03\t5\t0\nbuffer_add/add__full_buffer__cap_10000_\twall_time\t1.227909e-08\t1.224626e-08\t1.230893e-08\t2.552074e-03\t5\t0\nbuffer_create/create_fill_cap_100\talloc_words\t2.116000e+03\t2.116000e+03\t2.116000e+03\t0.000000e+00\t5\t0\nbuffer_create/create_fill_cap_100\tcpu_time\t1.430951e-06\t1.419749e-06\t1.441342e-06\t7.544736e-03\t5\t0\nbuffer_create/create_fill_cap_100\twall_time\t1.433376e-06\t1.418645e-06\t1.445267e-06\t9.286558e-03\t5\t1\nbuffer_create/create_fill_cap_1000\talloc_words\t2.101600e+04\t2.101600e+04\t2.101600e+04\t0.000000e+00\t5\t0\nbuffer_create/create_fill_cap_1000\tcpu_time\t1.815813e-05\t1.805756e-05\t1.829634e-05\t6.575071e-03\t5\t1\nbuffer_create/create_fill_cap_1000\twall_time\t1.817186e-05\t1.806923e-05\t1.829684e-05\t6.262546e-03\t5\t1\nbuffer_create/create_fill_cap_10000\talloc_words\t2.100160e+05\t2.100160e+05\t2.100160e+05\t0.000000e+00\t5\t0\nbuffer_create/create_fill_cap_10000\tcpu_time\t1.670111e-04\t1.660089e-04\t1.677894e-04\t5.330490e-03\t5\t0\nbuffer_create/create_fill_cap_10000\twall_time\t1.673292e-04\t1.663936e-04\t1.679940e-04\t4.782197e-03\t5\t1\ngae/compute_n_1024\talloc_words\t2.053000e+03\t2.053000e+03\t2.053000e+03\t0.000000e+00\t5\t0\ngae/compute_n_1024\tcpu_time\t4.432903e-06\t4.388681e-06\t4.472044e-06\t9.402727e-03\t5\t0\ngae/compute_n_1024\twall_time\t4.434699e-06\t4.389286e-06\t4.487822e-06\t1.110972e-02\t5\t0\ngae/compute_n_256\talloc_words\t5.170000e+02\t5.170000e+02\t5.170000e+02\t0.000000e+00\t5\t0\ngae/compute_n_256\tcpu_time\t8.899175e-07\t8.873495e-07\t8.928650e-07\t3.098850e-03\t5\t1\ngae/compute_n_256\twall_time\t8.917142e-07\t8.886773e-07\t8.955820e-07\t3.871563e-03\t5\t1\ngae/compute_n_4096\talloc_words\t8.197000e+03\t8.197000e+03\t8.197000e+03\t0.000000e+00\t5\t0\ngae/compute_n_4096\tcpu_time\t1.953902e-05\t1.938466e-05\t1.977415e-05\t9.967171e-03\t5\t0\ngae/compute_n_4096\twall_time\t1.954775e-05\t1.938586e-05\t1.974381e-05\t9.155682e-03\t5\t0\ngae_from_values/compute_from_values_n_1024\talloc_words\t5.130000e+03\t5.130000e+03\t5.130000e+03\t0.000000e+00\t5\t0\ngae_from_values/compute_from_values_n_1024\tcpu_time\t7.572262e-06\t7.451256e-06\t7.687885e-06\t1.562475e-02\t5\t0\ngae_from_values/compute_from_values_n_1024\twall_time\t7.577999e-06\t7.456862e-06\t7.698184e-06\t1.592253e-02\t5\t0\ngae_from_values/compute_from_values_n_256\talloc_words\t1.290000e+03\t1.290000e+03\t1.290000e+03\t0.000000e+00\t5\t0\ngae_from_values/compute_from_values_n_256\tcpu_time\t1.573065e-06\t1.567452e-06\t1.579053e-06\t3.687355e-03\t5\t0\ngae_from_values/compute_from_values_n_256\twall_time\t1.573980e-06\t1.567490e-06\t1.582763e-06\t4.851611e-03\t5\t0\ngae_from_values/compute_from_values_n_4096\talloc_words\t2.049000e+04\t2.049000e+04\t2.049000e+04\t0.000000e+00\t5\t0\ngae_from_values/compute_from_values_n_4096\tcpu_time\t3.301526e-05\t3.259093e-05\t3.336488e-05\t1.172117e-02\t5\t0\ngae_from_values/compute_from_values_n_4096\twall_time\t3.310795e-05\t3.264492e-05\t3.349818e-05\t1.288613e-02\t5\t0\nnormalize/normalize_n_1024\talloc_words\t3.083000e+03\t3.083000e+03\t3.083000e+03\t0.000000e+00\t5\t0\nnormalize/normalize_n_1024\tcpu_time\t8.499536e-06\t8.421020e-06\t8.595588e-06\t1.026924e-02\t5\t2\nnormalize/normalize_n_1024\twall_time\t8.507517e-06\t8.425891e-06\t8.596021e-06\t9.998763e-03\t5\t2\nnormalize/normalize_n_256\talloc_words\t7.790000e+02\t7.790000e+02\t7.790000e+02\t0.000000e+00\t5\t0\nnormalize/normalize_n_256\tcpu_time\t1.998406e-06\t1.990551e-06\t2.005793e-06\t3.813604e-03\t5\t0\nnormalize/normalize_n_256\twall_time\t1.999451e-06\t1.990430e-06\t2.006996e-06\t4.142583e-03\t5\t0\nnormalize/normalize_n_4096\talloc_words\t1.229900e+04\t1.229900e+04\t1.229900e+04\t0.000000e+00\t5\t0\nnormalize/normalize_n_4096\tcpu_time\t3.403318e-05\t3.376822e-05\t3.423778e-05\t6.898544e-03\t5\t0\nnormalize/normalize_n_4096\twall_time\t3.405293e-05\t3.381521e-05\t3.430319e-05\t7.164939e-03\t5\t0\nreturns/returns_n_1024\talloc_words\t1.025000e+03\t1.025000e+03\t1.025000e+03\t0.000000e+00\t5\t0\nreturns/returns_n_1024\tcpu_time\t2.979521e-06\t2.946996e-06\t3.001828e-06\t9.201606e-03\t5\t0\nreturns/returns_n_1024\twall_time\t2.994462e-06\t2.963207e-06\t3.022511e-06\t9.902417e-03\t5\t0\nreturns/returns_n_256\talloc_words\t2.570000e+02\t2.570000e+02\t2.570000e+02\t0.000000e+00\t5\t0\nreturns/returns_n_256\tcpu_time\t6.113050e-07\t6.027021e-07\t6.163104e-07\t1.113047e-02\t5\t1\nreturns/returns_n_256\twall_time\t6.123541e-07\t6.072009e-07\t6.175059e-07\t8.414242e-03\t5\t1\nreturns/returns_n_4096\talloc_words\t4.097000e+03\t4.097000e+03\t4.097000e+03\t0.000000e+00\t5\t0\nreturns/returns_n_4096\tcpu_time\t1.194814e-05\t1.181334e-05\t1.212174e-05\t1.290600e-02\t5\t0\nreturns/returns_n_4096\twall_time\t1.196891e-05\t1.180767e-05\t1.211994e-05\t1.304521e-02\t5\t0\n"
},
{
"path": "packages/fehu/doc/01-getting-started.md",
"content": "# Getting Started\n\nThis guide covers the basics: creating environments, running the step loop,\nunderstanding spaces, and using the built-in environments.\n\n## Installation\n\n\n```bash\nopam install fehu\n```\n\nOr build from source:\n\n\n```bash\ngit clone https://github.com/raven-ml/raven\ncd raven && dune build fehu\n```\n\n## Creating an Environment\n\nEnvironments are created via factory functions in `Fehu_envs`. Randomness is\nprovided by the implicit RNG scope from `Nx.Rng.run`:\n\n```ocaml\nopen Fehu\n\nlet () = Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make () in\n ignore env\n```\n\nThe seed controls all randomness in the scope. Use the same seed to get\nthe same episode sequence.\n\n## The Step Loop\n\nAn environment follows a strict lifecycle: `reset` must be called before the\nfirst `step`, and again after any terminal step (terminated or truncated).\n\n```ocaml\nopen Fehu\n\nlet () = Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make () in\n\n (* Reset returns the initial observation and info *)\n let _obs, _info = Env.reset env () in\n\n (* Step returns observation, reward, terminated, truncated, info *)\n let s = Env.step env (Space.Discrete.of_int 0) in\n Printf.printf \"reward: %.1f, terminated: %b, truncated: %b\\n\"\n s.reward s.terminated s.truncated\n```\n\nA complete episode loop:\n\n```ocaml\nopen Fehu\n\nlet run_episode env =\n let _obs, _info = Env.reset env () in\n let done_ = ref false in\n let total_reward = ref 0.0 in\n while not !done_ do\n let act = Space.sample (Env.action_space env) in\n let s = Env.step env act in\n total_reward := !total_reward +. s.reward;\n done_ := s.terminated || s.truncated\n done;\n !total_reward\n\nlet () = Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make () in\n let _reward = run_episode env in ()\n```\n\n## Spaces\n\nSpaces define the valid observations and actions for an environment. They\nprovide sampling, validation, and serialization.\n\n### Discrete\n\nInteger choices. Used for environments with a finite number of actions (e.g.,\nleft/right).\n\n```ocaml\nopen Fehu\n\nlet space = Space.Discrete.create 4 (* actions 0, 1, 2, 3 *)\nlet _n = Space.Discrete.n space (* 4 *)\n\n(* Sample a random action (requires an Nx.Rng scope) *)\nlet _act = Nx.Rng.run ~seed:0 @@ fun () ->\n Space.sample space\n\n(* Convert between int and discrete element *)\nlet act = Space.Discrete.of_int 2\nlet _i = Space.Discrete.to_int act (* 2 *)\n\n(* Check membership *)\nlet _valid = Space.contains space act (* true *)\n```\n\n### Box\n\nContinuous vectors with per-dimension bounds. Used for continuous observations\n(e.g., position, velocity) and continuous actions.\n\n```ocaml\nopen Fehu\n\nlet space = Space.Box.create\n ~low:[| -1.0; -2.0 |]\n ~high:[| 1.0; 2.0 |]\n\nlet _low, _high = Space.Box.bounds space\nlet _obs = Nx.Rng.run ~seed:0 @@ fun () -> Space.sample space\n```\n\n### Other Space Types\n\n- **Multi_binary**: binary vectors of fixed length (multi-label scenarios)\n- **Multi_discrete**: multiple discrete axes with independent cardinalities\n- **Tuple**: fixed-length heterogeneous sequences\n- **Dict**: named fields with different space types\n- **Sequence**: variable-length homogeneous sequences\n- **Text**: character strings from a fixed alphabet\n\nAll spaces support `contains`, `sample`, `pack`/`unpack` (to/from the\nuniversal `Value.t` type), and `boundary_values`.\n\n## Available Environments\n\n### CartPole\n\nClassic cart-pole balancing. Push a cart left or right to keep a pole upright.\nReward is +1.0 per step. Terminates when the pole exceeds +/-12 degrees or the\ncart leaves +/-2.4. Truncates at 500 steps.\n\n- **Observation**: Box [4] -- x, x_dot, theta, theta_dot\n- **Actions**: Discrete 2 -- 0 = push left, 1 = push right\n\n```ocaml\nlet _env = Nx.Rng.run ~seed:42 @@ fun () -> Fehu_envs.Cartpole.make ()\n```\n\n### MountainCar\n\nA car in a valley must build momentum to climb a hill. Reward is -1.0 per\nstep. Terminates when position >= 0.5 with non-negative velocity. Truncates at\n200 steps.\n\n- **Observation**: Box [2] -- position, velocity\n- **Actions**: Discrete 3 -- 0 = push left, 1 = coast, 2 = push right\n\n```ocaml\nlet _env = Nx.Rng.run ~seed:42 @@ fun () -> Fehu_envs.Mountain_car.make ()\n```\n\n### GridWorld\n\n5x5 grid navigation with an obstacle. Agent starts at (0,0), goal at (4,4),\nobstacle at (2,2). Reward is +10.0 at goal, -1.0 otherwise. Truncates at 200\nsteps.\n\n- **Observation**: Multi_discrete [5; 5] -- (row, col)\n- **Actions**: Discrete 4 -- 0 = up, 1 = down, 2 = left, 3 = right\n\n```ocaml\nlet _env = Nx.Rng.run ~seed:42 @@ fun () -> Fehu_envs.Grid_world.make ()\n```\n\n### RandomWalk\n\nOne-dimensional random walk on [-10, 10]. Reward is -|position|. Terminates at\nboundaries or after 200 steps.\n\n- **Observation**: Box [1] in [-10.0, 10.0]\n- **Actions**: Discrete 2 -- 0 = left, 1 = right\n\n```ocaml\nlet _env = Nx.Rng.run ~seed:42 @@ fun () -> Fehu_envs.Random_walk.make ()\n```\n\n## Render Modes\n\nEnvironments can optionally render their state. Pass `~render_mode` when\ncreating the environment:\n\n```ocaml\nopen Fehu\n\nlet () = Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make\n ~render_mode:`Ansi () in\n\n let _obs, _info = Env.reset env () in\n let _s = Env.step env (Space.Discrete.of_int 0) in\n\n (* Render after reset or step *)\n match Env.render env with\n | Some text -> print_endline text\n | None -> ()\n```\n\nSupported render modes vary by environment: `Ansi` for text output,\n`Rgb_array` for pixel frames, `Human` for interactive display.\n\n## Next Steps\n\n- [Environments and Wrappers](../02-environments/) -- custom environments, wrappers, rendering, vectorized environments\n- [Collection and Evaluation](../03-collection-and-evaluation/) -- trajectory collection, replay buffers, GAE, evaluation\n"
},
{
"path": "packages/fehu/doc/02-environments.md",
"content": "# Environments and Wrappers\n\nThis guide covers creating custom environments, composing wrappers, rendering,\nand running vectorized environments.\n\n## The Env.t Type\n\nAn environment `('obs, 'act, 'render) Env.t` is parameterized by its\nobservation type, action type, and render type. The type system ensures that\npolicies, wrappers, and collection utilities all agree on these types.\n\nThe lifecycle is strict:\n\n1. Call `Env.reset` to get the initial observation\n2. Call `Env.step` with an action to advance one timestep\n3. When `terminated` or `truncated` is true, call `Env.reset` again\n4. Call `Env.close` when done (optional, releases resources)\n\nCalling `step` before `reset`, or after a terminal step without resetting,\nraises `Invalid_argument`.\n\n## Creating Custom Environments\n\nUse `Env.create` to build an environment from `reset` and `step` functions.\nBoth receive the environment handle as their first argument, which provides\naccess to spaces and lifecycle state. Random keys are drawn from the implicit\nRNG scope (see below).\n\n\n```ocaml\nopen Fehu\n\n(* A simple counting environment: agent must choose action 1 *)\nlet make_counter () =\n let count = ref 0 in\n Env.create\n ~id:\"Counter-v0\"\n ~observation_space:(Space.Discrete.create 100)\n ~action_space:(Space.Discrete.create 2)\n ~reset:(fun _env ?options:_ () ->\n count := 0;\n Space.Discrete.of_int 0, Info.empty)\n ~step:(fun _env action ->\n let a = Space.Discrete.to_int action in\n if a = 1 then incr count else count := 0;\n let obs = Space.Discrete.of_int !count in\n let terminated = !count >= 10 in\n Env.step_result ~observation:obs\n ~reward:(if a = 1 then 1.0 else -1.0)\n ~terminated ())\n ()\n```\n\n### RNG Management\n\nEnvironments draw random keys from the implicit RNG scope established by\n`Nx.Rng.run`. Any call to `Space.sample` or other random operations inside\n`reset` and `step` callbacks will use this scope automatically:\n\n\n```ocaml\nlet make_noisy_env () =\n Env.create\n ~observation_space:(Space.Box.create\n ~low:[| 0.0 |] ~high:[| 1.0 |])\n ~action_space:(Space.Discrete.create 2)\n ~reset:(fun env ?options:_ () ->\n let obs = Space.sample (Env.observation_space env) in\n obs, Info.empty)\n ~step:(fun env _action ->\n let obs = Space.sample (Env.observation_space env) in\n Env.step_result ~observation:obs\n ~reward:1.0 ())\n ()\n```\n\n## Wrappers\n\nWrappers transform an environment's observations, actions, or rewards without\nmodifying the inner environment. They compose: wrap a wrapper to stack\ntransformations.\n\n### map_observation\n\nTransform observations from reset and step:\n\n\n```ocaml\nopen Fehu\n\n(* Normalize observations to [0, 1] *)\nlet env = Env.map_observation\n ~observation_space:(Space.Box.create\n ~low:[| 0.0; 0.0; 0.0; 0.0 |]\n ~high:[| 1.0; 1.0; 1.0; 1.0 |])\n ~f:(fun obs info ->\n (* obs is a float32 tensor, transform it *)\n let normalized = normalize_fn obs in\n normalized, info)\n env\n```\n\nThe function `f` receives both the observation and the info dictionary,\nreturning both. This allows wrappers to pass metadata through info.\n\n### map_action\n\nTransform actions before they reach the inner environment:\n\n\n```ocaml\n(* Remap discrete actions *)\nlet env = Env.map_action\n ~action_space:(Space.Discrete.create 3)\n ~f:(fun act ->\n (* Map from 3-action to 2-action space *)\n let i = Space.Discrete.to_int act in\n Space.Discrete.of_int (if i >= 2 then 1 else i))\n env\n```\n\n### map_reward\n\nTransform rewards after each step:\n\n\n```ocaml\n(* Scale rewards *)\nlet env = Env.map_reward\n ~f:(fun ~reward ~info -> reward *. 0.01, info)\n env\n```\n\n### clip_action\n\nClamp continuous actions to the action space bounds. The wrapper relaxes the\naction space to accept any float values, then clips before forwarding:\n\n\n```ocaml\n(* Works with Box action spaces *)\nlet env = Env.clip_action env\n```\n\n### clip_observation\n\nClamp observations to specified bounds:\n\n\n```ocaml\nlet env = Env.clip_observation\n ~low:[| -1.0; -1.0 |]\n ~high:[| 1.0; 1.0 |]\n env\n```\n\n### time_limit\n\nEnforce a maximum episode length. When the limit is reached, the step's\n`truncated` flag is set to true:\n\n\n```ocaml\nlet env = Env.time_limit ~max_episode_steps:200 env\n```\n\n### Custom Wrappers with Env.wrap\n\nFor transformations that need full control over reset and step, use `Env.wrap`.\nThe wrapper shares the inner environment's lifecycle (RNG, closed flag, reset\nflag):\n\n\n```ocaml\nopen Fehu\n\n(* A wrapper that tracks episode reward *)\nlet with_episode_reward env =\n let episode_reward = ref 0.0 in\n Env.wrap\n ~observation_space:(Env.observation_space env)\n ~action_space:(Env.action_space env)\n ~reset:(fun inner ?options () ->\n episode_reward := 0.0;\n Env.reset inner ?options ())\n ~step:(fun inner action ->\n let s = Env.step inner action in\n episode_reward := !episode_reward +. s.reward;\n let info =\n if s.terminated || s.truncated then\n Info.set \"episode_reward\"\n (Info.float !episode_reward) s.info\n else s.info\n in\n { s with info })\n env\n```\n\n## Rendering\n\nEnvironments support optional rendering via render modes. Pass\n`~render_mode` at creation time:\n\n\n```ocaml\nlet env = Fehu_envs.Grid_world.make\n ~render_mode:`Ansi ()\n\nlet _obs, _info = Env.reset env ()\nmatch Env.render env with\n| Some (Text s) -> print_endline s\n| _ -> ()\n```\n\n### Render Rollout\n\n`Render.rollout` runs a policy and feeds rendered frames to a sink function:\n\n\n```ocaml\nopen Fehu\n\n(* Collect rendered frames from a policy rollout *)\nlet frames = ref [] in\nRender.rollout env\n ~policy:(fun _obs -> Space.sample (Env.action_space env))\n ~steps:100\n ~sink:(fun img -> frames := img :: !frames)\n ()\n```\n\n### Recording with on_render\n\n`Render.on_render` wraps an environment so that every frame after reset and\nstep is passed to a sink:\n\n\n```ocaml\nlet env = Render.on_render\n ~sink:(fun img -> save_frame img)\n env\n```\n\n## Vectorized Environments\n\n`Vec_env` runs multiple environment instances with batched inputs and outputs.\nTerminated or truncated episodes are automatically reset.\n\n\n```ocaml\nopen Fehu\n\nlet () = Nx.Rng.run ~seed:42 @@ fun () ->\n (* Create 4 parallel environments *)\n let envs = List.init 4 (fun _ -> Fehu_envs.Cartpole.make ()) in\n let vec = Vec_env.create envs in\n\n let n = Vec_env.num_envs vec in (* 4 *)\n\n (* Reset all environments *)\n let _observations, _infos = Vec_env.reset vec () in\n\n (* Step all environments with an array of actions *)\n let actions = Array.init n (fun _ -> Space.Discrete.of_int 0) in\n let _s = Vec_env.step vec actions in\n\n (* _s.observations, _s.rewards, _s.terminated, _s.truncated *)\n\n (* Clean up *)\n Vec_env.close vec\n```\n\nAll environments must have structurally identical observation and action spaces\n(checked via `Space.equal_spec`). On terminal steps, the original terminal\nobservation is stored in the step info under `\"final_observation\"` as a packed\n`Value.t`, and the terminal info under `\"final_info\"`.\n\n## Next Steps\n\n- [Getting Started](../01-getting-started/) -- installation, environments, spaces, step loop\n- [Collection and Evaluation](../03-collection-and-evaluation/) -- trajectory collection, replay buffers, GAE, evaluation\n"
},
{
"path": "packages/fehu/doc/03-collection-and-evaluation.md",
"content": "# Collection, Buffers, and Evaluation\n\nThis guide covers trajectory collection, replay buffers, generalized advantage\nestimation, and policy evaluation.\n\n## Trajectory Collection\n\n`Collect` gathers agent-environment interactions into structure-of-arrays form\nfor batch processing.\n\n### Rollout\n\n`Collect.rollout` collects a fixed number of transitions. It resets the\nenvironment at the start and automatically on episode boundaries:\n\n\n```ocaml\nopen Fehu\n\nlet () = Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make () in\n\n (* The policy receives an observation and returns\n (action, log_prob option, value_estimate option) *)\n let policy _obs =\n let act = Space.sample (Env.action_space env) in\n (act, None, None)\n in\n\n let _trajectory = Collect.rollout env ~policy ~n_steps:1024 in ()\n```\n\nThe returned trajectory contains parallel arrays:\n\n\n```ocaml\nlet n = Collect.length trajectory (* 1024 *)\nlet obs = trajectory.observations (* 'obs array *)\nlet acts = trajectory.actions (* 'act array *)\nlet rews = trajectory.rewards (* float array *)\nlet next_obs = trajectory.next_observations (* 'obs array *)\nlet terms = trajectory.terminated (* bool array *)\nlet truncs = trajectory.truncated (* bool array *)\nlet infos = trajectory.infos (* Info.t array *)\nlet log_ps = trajectory.log_probs (* float array option *)\nlet vals = trajectory.values (* float array option *)\n```\n\nWhen the policy returns `Some log_prob` or `Some value`, those are collected\ninto `log_probs` and `values`. When any return is `None`, the corresponding\nfield is `None` for the entire trajectory.\n\n### Policy Signature\n\nThe policy function has the signature:\n\n```\n'obs -> 'act * float option * float option\n```\n\nThe three components are:\n\n1. **action**: the action to take\n2. **log_prob** (optional): the log-probability of the action under the current policy, used for importance sampling in PPO\n3. **value** (optional): the estimated value of the current state, used for GAE computation\n\nFor a simple random policy, return `None` for both:\n\n\n```ocaml\nlet random_policy _obs =\n let act = Space.sample (Env.action_space env) in\n (act, None, None)\n```\n\nFor a neural network policy with value head:\n\n\n```ocaml\nlet nn_policy obs =\n let logits, value = forward_pass model obs in\n let act = sample_from_logits logits in\n let log_prob = log_prob_of logits act in\n (act, Some log_prob, Some value)\n```\n\n### Episodes\n\n`Collect.episodes` collects complete episodes, one trajectory per episode:\n\n\n```ocaml\nlet episodes = Collect.episodes env\n ~policy ~n_episodes:10\n ~max_steps:500 ()\n\n(* episodes is a ('obs, 'act) Collect.t list *)\nlet total_rewards = List.map (fun traj ->\n Array.fold_left (+.) 0.0 traj.rewards) episodes\n```\n\nEach episode runs until termination, truncation, or `max_steps` (default\n1000).\n\n### Concatenating Trajectories\n\n`Collect.concat` merges multiple trajectories into one:\n\n\n```ocaml\nlet combined = Collect.concat [traj1; traj2; traj3]\n```\n\nOptional fields (`log_probs`, `values`) are kept only if present in all inputs.\n\n## Replay Buffers\n\n`Buffer` provides a fixed-capacity circular buffer for off-policy experience\nstorage. It stores individual transitions and supports uniform random sampling.\n\n### Creating and Filling\n\n\n```ocaml\nopen Fehu\n\nlet buf = Buffer.create ~capacity:10_000\n\n(* Add transitions one at a time *)\nBuffer.add buf {\n observation = obs;\n action = act;\n reward = 1.0;\n next_observation = next_obs;\n terminated = false;\n truncated = false;\n}\n\nlet n = Buffer.size buf (* number of stored transitions *)\nlet full = Buffer.is_full buf (* true when at capacity *)\nlet cap = Buffer.capacity buf (* 10000 *)\n```\n\nWhen the buffer is full, new transitions overwrite the oldest ones.\n\n### Sampling\n\nDraw a batch of transitions uniformly at random (with replacement):\n\n\n```ocaml\nlet batch = Nx.Rng.run ~seed:0 @@ fun () ->\n Buffer.sample buf ~batch_size:64\n\n(* batch is a transition array *)\nlet _obs_0 = batch.(0).observation\nlet _rew_0 = batch.(0).reward\n```\n\nFor structure-of-arrays form (more convenient for training):\n\n\n```ocaml\nlet (observations, actions, rewards,\n next_observations, terminated, truncated) =\n Nx.Rng.run ~seed:0 @@ fun () ->\n Buffer.sample_arrays buf ~batch_size:64\n```\n\n### Clearing\n\n\n```ocaml\nBuffer.clear buf (* removes all transitions, keeps storage allocated *)\n```\n\n## Generalized Advantage Estimation\n\n`Gae` computes advantages and returns for policy gradient methods. It correctly\nhandles the distinction between terminated and truncated episodes:\n\n- **Terminated**: the episode ended naturally (e.g., pole fell). Bootstrap\n value is zero.\n- **Truncated**: the episode was cut short (e.g., time limit). Bootstrap value\n comes from `next_values`.\n\n### Computing Advantages\n\n\n```ocaml\nopen Fehu\n\n(* From a trajectory with value estimates *)\nlet advantages, returns = Gae.compute\n ~rewards:trajectory.rewards\n ~values:(Option.get trajectory.values)\n ~terminated:trajectory.terminated\n ~truncated:trajectory.truncated\n ~next_values (* V(s_{t+1}) for each step *)\n ~gamma:0.99 (* discount factor *)\n ~lambda:0.95 (* GAE smoothing parameter *)\n```\n\nWhen you have values from a value network and the last value estimate,\n`compute_from_values` builds `next_values` for you:\n\n\n```ocaml\nlet advantages, returns = Gae.compute_from_values\n ~rewards:trajectory.rewards\n ~values:(Option.get trajectory.values)\n ~terminated:trajectory.terminated\n ~truncated:trajectory.truncated\n ~last_value:0.0 (* V(s_T) for the final state *)\n ~gamma:0.99\n ~lambda:0.95\n```\n\n### Monte Carlo Returns\n\nFor simpler algorithms that do not need advantages:\n\n\n```ocaml\nlet rets = Gae.returns\n ~rewards:trajectory.rewards\n ~terminated:trajectory.terminated\n ~truncated:trajectory.truncated\n ~gamma:0.99\n```\n\n### Normalizing Advantages\n\nNormalize to zero mean and unit variance for training stability:\n\n\n```ocaml\nlet normalized = Gae.normalize advantages\n(* or with custom epsilon *)\nlet normalized = Gae.normalize ~eps:1e-6 advantages\n```\n\n## Policy Evaluation\n\n`Eval.run` runs a deterministic or stochastic policy over multiple episodes\nand reports summary statistics:\n\n\n```ocaml\nopen Fehu\n\nlet () = Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make () in\n\n (* Evaluate a random policy *)\n let stats = Eval.run env\n ~policy:(fun _obs -> Space.sample (Env.action_space env))\n ~n_episodes:100\n ~max_steps:500\n ()\n in\n Printf.printf\n \"Episodes: %d, Mean reward: %.1f +/- %.1f, Mean length: %.0f\\n\"\n stats.n_episodes\n stats.mean_reward\n stats.std_reward\n stats.mean_length\n```\n\nThe evaluation policy has a simpler signature than the collection policy: it\nonly returns an action, not log-probs or value estimates:\n\n```\n'obs -> 'act\n```\n\n`Eval.run` resets the environment between episodes. Default `n_episodes` is 10\nand default `max_steps` is 1000.\n\n## Putting It Together\n\nA typical PPO-style training iteration using these utilities:\n\n\n```ocaml\nopen Fehu\n\n(* 1. Collect rollout *)\nlet trajectory = Collect.rollout env\n ~policy:(fun obs ->\n let act, log_prob, value = nn_policy obs in\n (act, Some log_prob, Some value))\n ~n_steps:2048\n\n(* 2. Compute advantages *)\nlet last_value = estimate_value model last_obs in\nlet advantages, returns = Gae.compute_from_values\n ~rewards:trajectory.rewards\n ~values:(Option.get trajectory.values)\n ~terminated:trajectory.terminated\n ~truncated:trajectory.truncated\n ~last_value\n ~gamma:0.99 ~lambda:0.95\n\nlet advantages = Gae.normalize advantages\n\n(* 3. Update policy using trajectory data + advantages *)\n(* ... your PPO update here ... *)\n\n(* 4. Evaluate *)\nlet stats = Eval.run env\n ~policy:(fun obs -> greedy_action model obs)\n ~n_episodes:10 ()\n```\n\n## Next Steps\n\n- [Getting Started](../01-getting-started/) -- installation, environments, spaces, step loop\n- [Environments and Wrappers](../02-environments/) -- custom environments, wrappers, rendering, vectorized environments\n"
},
{
"path": "packages/fehu/doc/04-gymnasium-comparison.md",
"content": "# Fehu vs. Gymnasium -- A Practical Comparison\n\nThis guide explains how Fehu's reinforcement learning API relates to Python's [Gymnasium](https://gymnasium.farama.org/) (and [Stable Baselines3](https://stable-baselines3.readthedocs.io/) for collection/buffer/GAE), focusing on:\n\n* How core concepts map (Env, Space, step loop, wrappers)\n* Where the APIs feel similar vs. deliberately different\n* How to translate common Gymnasium patterns into Fehu\n\nIf you already use Gymnasium, this should be enough to become productive in Fehu quickly.\n\n---\n\n## 1. Big-Picture Differences\n\n| Aspect | Gymnasium (Python) | Fehu (OCaml) |\n| --------------------- | -------------------------------------------------- | -------------------------------------------------------------------- |\n| Language | Dynamic, interpreted | Statically typed, compiled |\n| Environment type | `gymnasium.Env` | `('obs, 'act, 'render) Env.t` |\n| Observation/action | Untyped (`np.ndarray`, `int`, etc.) | Parametric: `'obs` and `'act` tracked in the type |\n| Spaces | `gymnasium.spaces.*` | `'a Space.t` with typed modules (`Space.Discrete`, `Space.Box`, ...) |\n| Step result | Tuple `(obs, reward, terminated, truncated, info)` | Record `Env.step` with named fields |\n| Wrappers | Subclassing `gymnasium.Wrapper` | `Env.wrap` or composable combinators (`map_observation`, etc.) |\n| Vectorized envs | `gymnasium.vector.SyncVectorEnv` | `Vec_env.create` |\n| Trajectory collection | External (Stable Baselines3, TorchRL) | Built-in: `Collect.rollout`, `Collect.episodes` |\n| Replay buffers | External (Stable Baselines3, TorchRL) | Built-in: `Buffer.create`, `Buffer.add`, `Buffer.sample` |\n| GAE | External (Stable Baselines3) | Built-in: `Gae.compute`, `Gae.returns`, `Gae.normalize` |\n| Policy evaluation | Manual loop or SB3 `evaluate_policy` | Built-in: `Eval.run` |\n| RNG | `np.random` / seed passed to `env.reset(seed=...)` | Implicit scope via `Nx.Rng.run ~seed` |\n| Rendering | String mode `\"human\"`, `\"rgb_array\"` | Polymorphic variants `` `Human ``, `` `Rgb_array ``, etc. |\n| Mutability | Environments are mutable objects | Environments are immutable handles; state is internal |\n\n**Fehu semantics to know (read once):**\n- `Env.reset` must be called before `Env.step`. After a terminal step, another `reset` is required.\n- Spaces validate observations and actions automatically -- `Env.step` raises if an action is outside the action space.\n- RNG is scoped: wrap your code in `Nx.Rng.run ~seed:42 (fun () -> ...)` instead of passing seeds to individual calls.\n- Trajectory collection, replay buffers, GAE, and evaluation are built into Fehu, not external libraries.\n\n---\n\n## 2. Spaces\n\n### 2.1 Discrete\n\n**Gymnasium**\n\n```python\nimport gymnasium as gym\n\nspace = gym.spaces.Discrete(5) # {0, 1, 2, 3, 4}\nspace = gym.spaces.Discrete(5, start=1) # {1, 2, 3, 4, 5}\n\nsample = space.sample()\nassert space.contains(sample)\n```\n\n**Fehu**\n\n\n```ocaml\nopen Fehu\n\nlet space = Space.Discrete.create 5 (* {0, 1, 2, 3, 4} *)\nlet space = Space.Discrete.create ~start:1 5 (* {1, 2, 3, 4, 5} *)\n\nlet sample = Space.sample space\nlet valid = Space.contains space sample\n\nlet n = Space.Discrete.n space (* 5 *)\nlet start = Space.Discrete.start space (* 1 *)\n\n(* Convert between discrete elements and ints *)\nlet action = Space.Discrete.of_int 3\nlet value = Space.Discrete.to_int action\n```\n\nDiscrete elements are `(int32, Nx.int32_elt) Nx.t` scalars, not bare OCaml ints.\n\n### 2.2 Box (continuous)\n\n**Gymnasium**\n\n```python\nimport numpy as np\n\nspace = gym.spaces.Box(\n low=np.array([-1.0, -2.0]),\n high=np.array([1.0, 2.0]),\n dtype=np.float32,\n)\nsample = space.sample()\n```\n\n**Fehu**\n\n\n```ocaml\nlet space =\n Space.Box.create\n ~low:[| -1.0; -2.0 |]\n ~high:[| 1.0; 2.0 |]\n\nlet sample = Space.sample space\nlet (low, high) = Space.Box.bounds space\n```\n\nBox elements are `(float, Nx.float32_elt) Nx.t` tensors. Infinite bounds are allowed; sampling falls back to uniform draws in `[-1e6, 1e6]` clamped to bounds.\n\n### 2.3 Multi_binary\n\n**Gymnasium**\n\n```python\nspace = gym.spaces.MultiBinary(4) # {0,1}^4\n```\n\n**Fehu**\n\n\n```ocaml\nlet space = Space.Multi_binary.create 4\n```\n\nElements are `(int32, Nx.int32_elt) Nx.t` vectors with values 0 or 1.\n\n### 2.4 Multi_discrete\n\n**Gymnasium**\n\n```python\nspace = gym.spaces.MultiDiscrete([3, 5, 2]) # 3 axes: {0..2}, {0..4}, {0..1}\n```\n\n**Fehu**\n\n\n```ocaml\nlet space = Space.Multi_discrete.create [| 3; 5; 2 |]\n```\n\n### 2.5 Composite spaces\n\n**Gymnasium**\n\n```python\nspace = gym.spaces.Tuple((\n gym.spaces.Discrete(3),\n gym.spaces.Box(low=0.0, high=1.0, shape=(2,)),\n))\n\nspace = gym.spaces.Dict({\n \"position\": gym.spaces.Box(low=-10.0, high=10.0, shape=(3,)),\n \"velocity\": gym.spaces.Box(low=-1.0, high=1.0, shape=(3,)),\n})\n```\n\n**Fehu**\n\n\n```ocaml\nlet space =\n Space.Tuple.create [\n Space.Pack (Space.Discrete.create 3);\n Space.Pack (Space.Box.create ~low:[| 0.0; 0.0 |] ~high:[| 1.0; 1.0 |]);\n ]\n\nlet space =\n Space.Dict.create [\n (\"position\", Space.Pack (Space.Box.create ~low:[| -10.; -10.; -10. |] ~high:[| 10.; 10.; 10. |]));\n (\"velocity\", Space.Pack (Space.Box.create ~low:[| -1.; -1.; -1. |] ~high:[| 1.; 1.; 1. |]));\n ]\n```\n\nComposite space elements use `Value.t` for heterogeneous data: `Tuple.element = Value.t list`, `Dict.element = (string * Value.t) list`.\n\n### 2.6 Sequence and Text\n\n**Gymnasium**\n\n```python\nspace = gym.spaces.Sequence(gym.spaces.Discrete(5), seed=42)\nspace = gym.spaces.Text(max_length=32, charset=\"abcdef\")\n```\n\n**Fehu**\n\n\n```ocaml\nlet space = Space.Sequence.create ~max_length:10 (Space.Discrete.create 5)\nlet space = Space.Text.create ~charset:\"abcdef\" ~max_length:32 ()\n```\n\n### 2.7 Common operations\n\nAll space types share the same interface:\n\n\n```ocaml\nlet sample = Space.sample space (* random element *)\nlet valid = Space.contains space sample (* membership test *)\nlet spec = Space.spec space (* structural description *)\nlet shape = Space.shape space (* dimensionality, if defined *)\n\n(* Serialization via Value.t *)\nlet packed = Space.pack space sample\nlet unpacked = Space.unpack space packed (* (element, string) result *)\n\n(* Edge cases for testing *)\nlet edges = Space.boundary_values space\n```\n\n---\n\n## 3. Creating Environments\n\n### 3.1 From a registry\n\n**Gymnasium**\n\n```python\nenv = gym.make(\"CartPole-v1\", render_mode=\"human\")\n```\n\n**Fehu** does not have a global registry. Environments are constructed directly:\n\n\n```ocaml\nlet env =\n Env.create\n ~id:\"CartPole-v1\"\n ~observation_space:(Space.Box.create\n ~low:[| -4.8; Float.neg_infinity; -0.418; Float.neg_infinity |]\n ~high:[| 4.8; Float.infinity; 0.418; Float.infinity |])\n ~action_space:(Space.Discrete.create 2)\n ~render_mode:`Human\n ~render_modes:[\"human\"; \"rgb_array\"]\n ~reset:(fun _env ?options:_ () ->\n let obs = (* initial state *) in\n (obs, Info.empty))\n ~step:(fun _env action ->\n let obs = (* next state *) in\n Env.step_result ~observation:obs ~reward:1.0 ())\n ()\n```\n\n`Env.create` takes the observation space, action space, and two callbacks: `reset` and `step`. Optional `render` and `close` callbacks handle visualization and cleanup.\n\n### 3.2 Step result construction\n\n**Gymnasium** returns a flat tuple from `env.step()`:\n\n```python\nobs, reward, terminated, truncated, info = env.step(action)\n```\n\n**Fehu** uses a record with named fields, and provides a convenience constructor with defaults:\n\n\n```ocaml\n(* Inside a step callback *)\nEnv.step_result\n ~observation:obs\n ~reward:1.0\n ~terminated:false\n ~truncated:false\n ~info:Info.empty\n ()\n\n(* Defaults: reward=0., terminated=false, truncated=false, info=Info.empty *)\nEnv.step_result ~observation:obs ()\n```\n\n---\n\n## 4. Step Loop\n\n### 4.1 Basic episode\n\n**Gymnasium**\n\n```python\nenv = gym.make(\"CartPole-v1\")\nobs, info = env.reset(seed=42)\n\ntotal_reward = 0.0\nwhile True:\n action = env.action_space.sample()\n obs, reward, terminated, truncated, info = env.step(action)\n total_reward += reward\n if terminated or truncated:\n break\n\nenv.close()\n```\n\n**Fehu**\n\n\n```ocaml\nlet () =\n Nx.Rng.run ~seed:42 (fun () ->\n let env = (* create environment *) in\n let (obs, _info) = Env.reset env () in\n let obs = ref obs in\n let total_reward = ref 0.0 in\n let done_ = ref false in\n while not !done_ do\n let action = Space.sample (Env.action_space env) in\n let step = Env.step env action in\n obs := step.observation;\n total_reward := !total_reward +. step.reward;\n done_ := step.terminated || step.truncated\n done;\n Env.close env)\n```\n\nKey differences:\n- RNG is scoped with `Nx.Rng.run ~seed:42` rather than passed to `reset`.\n- Step results are accessed by field name (`step.observation`, `step.reward`).\n- `Env.step` raises `Invalid_argument` if called without a prior `reset` or after a terminal step without resetting.\n\n### 4.2 Multiple episodes\n\n**Gymnasium**\n\n```python\nfor episode in range(10):\n obs, info = env.reset()\n done = False\n while not done:\n action = policy(obs)\n obs, reward, terminated, truncated, info = env.step(action)\n done = terminated or truncated\n```\n\n**Fehu** -- manual loop or use `Collect.episodes`:\n\n\n```ocaml\n(* Manual *)\nlet () =\n Nx.Rng.run ~seed:0 (fun () ->\n let env = (* create environment *) in\n for _ep = 0 to 9 do\n let (obs, _info) = Env.reset env () in\n let obs = ref obs in\n let done_ = ref false in\n while not !done_ do\n let action = policy !obs in\n let step = Env.step env action in\n obs := step.observation;\n done_ := step.terminated || step.truncated\n done\n done;\n Env.close env)\n\n(* Or use Collect.episodes directly *)\nlet trajs =\n Nx.Rng.run ~seed:0 (fun () ->\n let env = (* create environment *) in\n Collect.episodes env\n ~policy:(fun obs -> (policy obs, None, None))\n ~n_episodes:10 ())\n```\n\n---\n\n## 5. Wrappers\n\n### 5.1 Gymnasium approach: subclassing\n\n**Gymnasium**\n\n```python\nclass NormalizeObservation(gym.Wrapper):\n def __init__(self, env, mean, std):\n super().__init__(env)\n self.mean = mean\n self.std = std\n\n def observation(self, obs):\n return (obs - self.mean) / self.std\n\nenv = NormalizeObservation(env, mean=0.0, std=1.0)\n```\n\n### 5.2 Fehu approach: composable functions\n\n**Fehu** provides `Env.wrap` for full control and specialized combinators for common patterns.\n\n**`map_observation`** -- transform observations:\n\n\n```ocaml\nlet normalized_env =\n Env.map_observation\n ~observation_space:obs_space\n ~f:(fun obs _info ->\n let normalized = (* normalize obs *) in\n (normalized, Info.empty))\n env\n```\n\n**`map_action`** -- transform actions before passing to the inner env:\n\n\n```ocaml\nlet remapped_env =\n Env.map_action\n ~action_space:new_action_space\n ~f:(fun new_action -> (* convert to inner action *))\n env\n```\n\n**`map_reward`** -- transform rewards:\n\n\n```ocaml\nlet scaled_env =\n Env.map_reward\n ~f:(fun ~reward ~info -> (reward *. 0.1, info))\n env\n```\n\n**`clip_action`** -- clamp continuous actions to bounds:\n\n\n```ocaml\n(* Gymnasium *)\n(* from gymnasium.wrappers import ClipAction *)\n(* env = ClipAction(env) *)\n\n(* Fehu *)\nlet clipped_env = Env.clip_action env\n```\n\n**`clip_observation`** -- clamp observations:\n\n\n```ocaml\nlet clipped_env =\n Env.clip_observation\n ~low:[| -5.0; -5.0 |]\n ~high:[| 5.0; 5.0 |]\n env\n```\n\n**`time_limit`** -- enforce maximum episode length:\n\n\n```ocaml\n(* Gymnasium *)\n(* from gymnasium.wrappers import TimeLimit *)\n(* env = TimeLimit(env, max_episode_steps=200) *)\n\n(* Fehu *)\nlet limited_env = Env.time_limit ~max_episode_steps:200 env\n```\n\n### 5.3 Full custom wrapper with `Env.wrap`\n\nWhen the combinators are not enough, use `Env.wrap` directly:\n\n\n```ocaml\nlet custom_env =\n Env.wrap\n ~observation_space:new_obs_space\n ~action_space:new_act_space\n ~reset:(fun inner ?options () ->\n let (obs, info) = Env.reset inner ?options () in\n (transform_obs obs, info))\n ~step:(fun inner action ->\n let step = Env.step inner (transform_action action) in\n { step with observation = transform_obs step.observation })\n env\n```\n\n`Env.wrap` receives the inner environment as the first argument to `reset`, `step`, `render`, and `close`. Guards (closed check, needs-reset check, space validation) are enforced automatically.\n\n### 5.4 Composing wrappers\n\nWrappers compose by chaining:\n\n\n```ocaml\nlet env =\n base_env\n |> Env.time_limit ~max_episode_steps:500\n |> Env.clip_action\n |> Env.map_reward ~f:(fun ~reward ~info -> (reward *. 0.01, info))\n```\n\n---\n\n## 6. Vectorized Environments\n\n### 6.1 Synchronous vectorization\n\n**Gymnasium**\n\n```python\nenvs = gym.vector.SyncVectorEnv([\n lambda: gym.make(\"CartPole-v1\") for _ in range(4)\n])\n\nobs, infos = envs.reset()\nactions = envs.action_space.sample() # batch of 4 actions\nobs, rewards, terminated, truncated, infos = envs.step(actions)\nenvs.close()\n```\n\n**Fehu**\n\n\n```ocaml\nlet venv =\n Vec_env.create [env1; env2; env3; env4]\n\nlet n = Vec_env.num_envs venv (* 4 *)\n\nlet (observations, infos) = Vec_env.reset venv ()\nlet actions = Array.init n (fun _ -> Space.sample (Vec_env.action_space venv)) in\nlet step = Vec_env.step venv actions\n\n(* step.observations : 'obs array -- one per env *)\n(* step.rewards : float array -- one per env *)\n(* step.terminated : bool array -- one per env *)\n(* step.truncated : bool array -- one per env *)\n(* step.infos : Info.t array -- one per env *)\n\nVec_env.close venv\n```\n\nKey differences:\n- `Vec_env.create` takes a list of already-constructed environments. All must have structurally identical spaces.\n- Terminated or truncated environments are automatically reset. The terminal observation is stored in the step's info under `\"final_observation\"` (as a packed `Value.t`), and the terminal info under `\"final_info\"`.\n- The step result is a record with named arrays, not a tuple.\n\n---\n\n## 7. Trajectory Collection\n\n### 7.1 Fixed-step rollout\n\n**Gymnasium + Stable Baselines3**\n\n```python\nfrom stable_baselines3.common.buffers import RolloutBuffer\n\n# Manual loop or SB3 internals\nobs, _ = env.reset()\nfor step in range(2048):\n action, log_prob, value = policy(obs)\n obs, reward, terminated, truncated, info = env.step(action)\n buffer.add(obs, action, reward, ...)\n if terminated or truncated:\n obs, _ = env.reset()\n```\n\n**Fehu** -- built-in:\n\n\n```ocaml\nlet trajectory =\n Collect.rollout env\n ~policy:(fun obs ->\n let action = (* select action *) in\n let log_prob = (* optional log probability *) in\n let value = (* optional value estimate *) in\n (action, Some log_prob, Some value))\n ~n_steps:2048\n```\n\n`Collect.rollout` handles resets on episode boundaries automatically and returns a `Collect.t` record:\n\n\n```ocaml\n(* Collect.t fields: *)\ntrajectory.observations (* 'obs array *)\ntrajectory.actions (* 'act array *)\ntrajectory.rewards (* float array *)\ntrajectory.next_observations (* 'obs array *)\ntrajectory.terminated (* bool array *)\ntrajectory.truncated (* bool array *)\ntrajectory.infos (* Info.t array *)\ntrajectory.log_probs (* float array option *)\ntrajectory.values (* float array option *)\n\nlet n = Collect.length trajectory\n```\n\n### 7.2 Complete episodes\n\n**Gymnasium + manual**\n\n```python\nepisodes = []\nfor _ in range(10):\n obs, _ = env.reset()\n episode = []\n done = False\n while not done:\n action = policy(obs)\n next_obs, reward, terminated, truncated, info = env.step(action)\n episode.append((obs, action, reward, next_obs, terminated, truncated))\n obs = next_obs\n done = terminated or truncated\n episodes.append(episode)\n```\n\n**Fehu** -- built-in:\n\n\n```ocaml\nlet episode_list =\n Collect.episodes env\n ~policy:(fun obs -> (policy obs, None, None))\n ~n_episodes:10\n ~max_steps:1000\n ()\n(* episode_list : ('obs, 'act) Collect.t list *)\n```\n\nEach element is one episode as a `Collect.t`. Concatenate them with `Collect.concat`:\n\n\n```ocaml\nlet all_transitions = Collect.concat episode_list\n```\n\n---\n\n## 8. Replay Buffers\n\n### 8.1 Standard replay buffer\n\n**Stable Baselines3**\n\n```python\nfrom stable_baselines3.common.buffers import ReplayBuffer\n\nbuffer = ReplayBuffer(buffer_size=100_000, observation_space=..., action_space=...)\nbuffer.add(obs, next_obs, action, reward, done, infos)\nbatch = buffer.sample(batch_size=256)\n```\n\n**Fehu** -- built-in:\n\n\n```ocaml\nlet buf = Buffer.create ~capacity:100_000\n\nlet () =\n Buffer.add buf {\n Buffer.observation = obs;\n action;\n reward = 1.0;\n next_observation = next_obs;\n terminated = false;\n truncated = false;\n }\n\n(* Uniform random sampling *)\nlet batch = Buffer.sample buf ~batch_size:256\n(* batch : ('obs, 'act) Buffer.transition array *)\n\n(* Structure-of-arrays form for training loops *)\nlet (observations, actions, rewards, next_observations, terminated, truncated) =\n Buffer.sample_arrays buf ~batch_size:256\n```\n\n### 8.2 Buffer queries\n\n\n```ocaml\nlet n = Buffer.size buf (* current number of stored transitions *)\nlet cap = Buffer.capacity buf (* maximum capacity *)\nlet full = Buffer.is_full buf (* true when size = capacity *)\nlet () = Buffer.clear buf (* remove all transitions, keep storage *)\n```\n\n---\n\n## 9. GAE and Returns\n\n### 9.1 Generalized Advantage Estimation\n\n**Stable Baselines3** (internal)\n\n```python\n# SB3 computes GAE internally in on-policy algorithms\n# or manually:\nimport numpy as np\n\ndef compute_gae(rewards, values, dones, next_values, gamma=0.99, lam=0.95):\n advantages = np.zeros_like(rewards)\n last_gae = 0\n for t in reversed(range(len(rewards))):\n delta = rewards[t] + gamma * next_values[t] * (1 - dones[t]) - values[t]\n advantages[t] = last_gae = delta + gamma * lam * (1 - dones[t]) * last_gae\n returns = advantages + values\n return advantages, returns\n```\n\n**Fehu** -- built-in, with correct terminated/truncated handling:\n\n\n```ocaml\nlet (advantages, returns) =\n Gae.compute\n ~rewards:trajectory.rewards\n ~values:(Option.get trajectory.values)\n ~terminated:trajectory.terminated\n ~truncated:trajectory.truncated\n ~next_values (* float array: V(s_{t+1}) for each t *)\n ~gamma:0.99\n ~lambda:0.95\n```\n\nWhen you have values from a rollout and a final bootstrap value:\n\n\n```ocaml\nlet (advantages, returns) =\n Gae.compute_from_values\n ~rewards:trajectory.rewards\n ~values:(Option.get trajectory.values)\n ~terminated:trajectory.terminated\n ~truncated:trajectory.truncated\n ~last_value:0.0\n ~gamma:0.99\n ~lambda:0.95\n```\n\n`compute_from_values` builds `next_values` from `values` and `last_value` automatically: `next_values.(t) = values.(t+1)` for `t < n-1`, and `next_values.(n-1) = last_value`.\n\n### 9.2 Monte Carlo returns\n\n**Manual Python**\n\n```python\ndef discounted_returns(rewards, dones, gamma=0.99):\n returns = np.zeros_like(rewards)\n running = 0.0\n for t in reversed(range(len(rewards))):\n running = rewards[t] + gamma * running * (1 - dones[t])\n returns[t] = running\n return returns\n```\n\n**Fehu**\n\n\n```ocaml\nlet mc_returns =\n Gae.returns\n ~rewards:trajectory.rewards\n ~terminated:trajectory.terminated\n ~truncated:trajectory.truncated\n ~gamma:0.99\n```\n\n### 9.3 Normalization\n\n\n```ocaml\nlet normalized_advantages = Gae.normalize advantages\nlet normalized_custom = Gae.normalize ~eps:1e-5 advantages\n```\n\n---\n\n## 10. Policy Evaluation\n\n**Gymnasium + Stable Baselines3**\n\n```python\nfrom stable_baselines3.common.evaluation import evaluate_policy\n\nmean_reward, std_reward = evaluate_policy(\n model, env, n_eval_episodes=10, deterministic=True\n)\n```\n\n**Fehu** -- built-in:\n\n\n```ocaml\nlet stats =\n Eval.run env\n ~policy:(fun obs -> (* deterministic action *))\n ~n_episodes:10\n ~max_steps:1000\n ()\n\n(* stats.mean_reward : float *)\n(* stats.std_reward : float *)\n(* stats.mean_length : float *)\n(* stats.n_episodes : int *)\n```\n\n`Eval.run` resets the environment between episodes and collects total reward and episode length across all episodes.\n\n---\n\n## 11. Rendering\n\n### 11.1 Render modes\n\n**Gymnasium**\n\n```python\nenv = gym.make(\"CartPole-v1\", render_mode=\"human\")\nenv.reset()\nenv.step(action)\nframe = env.render() # None for \"human\", np.ndarray for \"rgb_array\"\n```\n\n**Fehu**\n\n\n```ocaml\nlet env =\n Env.create\n ~render_mode:`Human\n ~render_modes:[\"human\"; \"rgb_array\"]\n ~render:(fun () -> (* return 'render option *))\n (* ... *)\n ()\n\nlet frame = Env.render env (* 'render option *)\n```\n\nRender modes are polymorphic variants: `` `Human ``, `` `Rgb_array ``, `` `Ansi ``, `` `Svg ``, `` `Custom of string ``.\n\n### 11.2 Frame type\n\nFor `Rgb_array` environments, Fehu uses `Render.image`:\n\n\n```ocaml\n(* Render.image fields: *)\n(* width : int *)\n(* height : int *)\n(* pixel_format : Render.Pixel.format (Rgb|Rgba|Gray) *)\n(* data : uint8 bigarray *)\n```\n\n### 11.3 Recording rendered rollouts\n\n**Gymnasium**\n\n```python\nfrom gymnasium.wrappers import RecordVideo\nenv = RecordVideo(env, video_folder=\"./videos\")\n```\n\n**Fehu** -- use `Render.rollout` or `Render.on_render`:\n\n\n```ocaml\n(* Run a policy and feed frames to a sink *)\nRender.rollout env\n ~policy:(fun obs -> (* action *))\n ~steps:500\n ~sink:(fun frame -> (* save or display frame *))\n ()\n\n(* Or wrap the env to capture every rendered frame *)\nlet recording_env =\n Render.on_render\n ~sink:(fun frame -> (* process frame *))\n env\n```\n\n---\n\n## 12. Info Dictionaries\n\n**Gymnasium** uses plain Python dicts for info:\n\n```python\nobs, info = env.reset()\nprint(info.get(\"elapsed_steps\", 0))\n```\n\n**Fehu** uses typed `Info.t` dictionaries with `Value.t` values:\n\n\n```ocaml\nlet info = Info.of_list [\n (\"elapsed_steps\", Info.int 42);\n (\"success\", Info.bool true);\n]\n\nlet steps = Info.find \"elapsed_steps\" info (* Value.t option *)\nlet steps = Info.find_exn \"elapsed_steps\" info (* Value.t, raises on missing *)\n\nlet info' = Info.set \"custom_key\" (Info.float 3.14) info\nlet info' = Info.merge info1 info2 (* info2 wins on conflicts *)\nlet is_empty = Info.is_empty info\n```\n\n---\n\n## 13. Quick Cheat Sheet\n\n| Task | Gymnasium / SB3 | Fehu |\n| -------------------- | ------------------------------------------------- | --------------------------------------------------------------------------------- |\n| Create env | `gym.make(\"CartPole-v1\")` | `Env.create ~observation_space ~action_space ~reset ~step ()` |\n| Reset | `obs, info = env.reset(seed=42)` | `let (obs, info) = Env.reset env ()` |\n| Step | `obs, r, term, trunc, info = env.step(a)` | `let s = Env.step env a` (record fields) |\n| Close | `env.close()` | `Env.close env` |\n| Discrete space | `gym.spaces.Discrete(5)` | `Space.Discrete.create 5` |\n| Box space | `gym.spaces.Box(low, high)` | `Space.Box.create ~low ~high` |\n| Sample from space | `space.sample()` | `Space.sample space` |\n| Contains check | `space.contains(x)` | `Space.contains space x` |\n| Observation wrapper | `class W(gym.ObservationWrapper)` | `Env.map_observation ~observation_space ~f env` |\n| Action wrapper | `class W(gym.ActionWrapper)` | `Env.map_action ~action_space ~f env` |\n| Reward wrapper | `class W(gym.RewardWrapper)` | `Env.map_reward ~f env` |\n| Clip actions | `ClipAction(env)` | `Env.clip_action env` |\n| Time limit | `TimeLimit(env, max_episode_steps=N)` | `Env.time_limit ~max_episode_steps:N env` |\n| Vectorize | `gym.vector.SyncVectorEnv([...])` | `Vec_env.create [env1; env2; ...]` |\n| Rollout N steps | Manual loop / SB3 internal | `Collect.rollout env ~policy ~n_steps` |\n| Collect N episodes | Manual loop | `Collect.episodes env ~policy ~n_episodes ()` |\n| Replay buffer | `ReplayBuffer(buffer_size=N, ...)` | `Buffer.create ~capacity:N` |\n| Add to buffer | `buffer.add(obs, next_obs, ...)` | `Buffer.add buf transition` |\n| Sample from buffer | `buffer.sample(batch_size=B)` | `Buffer.sample buf ~batch_size:B` |\n| GAE | SB3 internal / manual | `Gae.compute ~rewards ~values ~terminated ~truncated ~next_values ~gamma ~lambda` |\n| Discounted returns | Manual loop | `Gae.returns ~rewards ~terminated ~truncated ~gamma` |\n| Normalize advantages | `(adv - mean) / std` | `Gae.normalize advantages` |\n| Evaluate policy | `evaluate_policy(model, env, n_eval_episodes=10)` | `Eval.run env ~policy ~n_episodes:10 ()` |\n| Render | `env.render()` | `Env.render env` |\n| Record frames | `RecordVideo(env, ...)` | `Render.on_render ~sink env` |\n| Seed RNG | `env.reset(seed=42)` | `Nx.Rng.run ~seed:42 (fun () -> ...)` |\n"
},
{
"path": "packages/fehu/doc/dune",
"content": "(mdx\n (files *.md)\n (package fehu)\n (libraries fehu fehu.envs nx))\n"
},
{
"path": "packages/fehu/doc/index.md",
"content": "# Fehu\n\nFehu is a reinforcement learning environment toolkit for OCaml. It provides\ntype-safe environments, composable wrappers, trajectory collection, replay\nbuffers, GAE computation, policy evaluation, and vectorized environments.\n\nFehu follows the Gymnasium interface pattern: environments expose `reset` and\n`step` with typed observation and action spaces. Wrappers compose freely.\nCollection and evaluation utilities handle the plumbing between environments\nand training loops.\n\n## Features\n\n- **Type-safe environments**: observation and action spaces are encoded in the type system\n- **Rich space types**: Discrete, Box, Multi_binary, Multi_discrete, Tuple, Dict, Sequence, Text\n- **Composable wrappers**: map_observation, map_action, map_reward, clip_action, clip_observation, time_limit\n- **Trajectory collection**: rollout and episode collection in structure-of-arrays form\n- **Replay buffers**: fixed-capacity circular buffer with uniform random sampling\n- **GAE**: generalized advantage estimation with proper terminated/truncated handling\n- **Policy evaluation**: run a policy over episodes and get mean/std reward statistics\n- **Vectorized environments**: run multiple environments with batched step and auto-reset\n- **Built-in environments**: CartPole, MountainCar, GridWorld, RandomWalk\n\n## Quick Start\n\nCreate an environment, run a random agent, and evaluate:\n\n```ocaml\nopen Fehu\n\nlet () = Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make () in\n\n (* Run one episode *)\n let _obs, _info = Env.reset env () in\n let done_ = ref false in\n let total_reward = ref 0.0 in\n while not !done_ do\n let act = Space.sample (Env.action_space env) in\n let s = Env.step env act in\n total_reward := !total_reward +. s.reward;\n done_ := s.terminated || s.truncated\n done;\n\n (* Evaluate over 10 episodes *)\n let _stats = Eval.run env\n ~policy:(fun _obs -> Space.sample (Env.action_space env))\n ~n_episodes:10 ()\n in ()\n```\n\n## Next Steps\n\n- [Getting Started](01-getting-started/) -- installation, environments, spaces, step loop\n- [Environments and Wrappers](02-environments/) -- custom environments, wrappers, rendering, vectorized environments\n- [Collection and Evaluation](03-collection-and-evaluation/) -- trajectory collection, replay buffers, GAE, evaluation\n"
},
{
"path": "packages/fehu/examples/01-random-agent/dune",
"content": "(executable\n (name main)\n (libraries nx rune fehu fehu.envs))\n"
},
{
"path": "packages/fehu/examples/01-random-agent/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* A random agent on CartPole-v1.\n\n Demonstrates the Env lifecycle: create, reset, step, render, close. Then uses\n Eval.run for batch evaluation. *)\n\nopen Fehu\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n Printf.printf \"Random Agent on CartPole-v1\\n\";\n Printf.printf \"===========================\\n\\n\";\n\n let env = Fehu_envs.Cartpole.make ~render_mode:`Ansi () in\n\n (* -- Manual episode loop ------------------------------------------------ *)\n Printf.printf \"Running 5 episodes with random actions...\\n\\n\";\n\n for episode = 1 to 5 do\n let obs = ref (fst (Env.reset env ())) in\n let total_reward = ref 0.0 in\n let steps = ref 0 in\n let done_ = ref false in\n while not !done_ do\n (* Show the first step of episode 1 *)\n (if episode = 1 && !steps = 0 then\n match Env.render env with\n | Some text -> Printf.printf \"%s\\n\" text\n | None -> ());\n let action = Space.sample (Env.action_space env) in\n let s = Env.step env action in\n total_reward := !total_reward +. s.reward;\n incr steps;\n obs := s.observation;\n done_ := s.terminated || s.truncated\n done;\n Printf.printf \" Episode %d: reward = %5.1f length = %3d\\n\" episode\n !total_reward !steps\n done;\n\n (* -- Batch evaluation with Eval.run ------------------------------------ *)\n Printf.printf \"\\nEvaluating over 100 episodes...\\n\\n\";\n\n let random_policy _obs = Space.sample (Env.action_space env) in\n let stats = Eval.run env ~policy:random_policy ~n_episodes:100 () in\n Printf.printf \" mean reward: %6.2f +/- %.2f\\n\" stats.mean_reward\n stats.std_reward;\n Printf.printf \" mean length: %6.1f\\n\" stats.mean_length;\n\n Env.close env;\n Printf.printf \"\\nDone.\\n\"\n"
},
{
"path": "packages/fehu/examples/02-q-learning/dune",
"content": "(executable\n (name main)\n (libraries nx rune fehu fehu.envs))\n"
},
{
"path": "packages/fehu/examples/02-q-learning/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Tabular Q-learning on CartPole-v1.\n\n Discretizes the continuous 4D observation into bins, learns a Q-table with\n epsilon-greedy exploration and temporal difference updates. Uses Eval.run for\n periodic evaluation. *)\n\nopen Fehu\n\n(* Hyperparameters *)\n\nlet n_bins = 12\nlet n_actions = 2\nlet alpha = 0.1\nlet gamma = 0.99\nlet epsilon_start = 1.0\nlet epsilon_end = 0.01\nlet epsilon_decay = 2000.0\nlet n_episodes = 10_000\nlet eval_interval = 500\n\n(* Sparkline *)\n\nlet sparkline values =\n let blocks =\n [|\n \"\\xe2\\x96\\x81\";\n \"\\xe2\\x96\\x82\";\n \"\\xe2\\x96\\x83\";\n \"\\xe2\\x96\\x84\";\n \"\\xe2\\x96\\x85\";\n \"\\xe2\\x96\\x86\";\n \"\\xe2\\x96\\x87\";\n \"\\xe2\\x96\\x88\";\n |]\n in\n let lo = Array.fold_left Float.min Float.infinity values in\n let hi = Array.fold_left Float.max Float.neg_infinity values in\n let range = hi -. lo in\n if range < 1e-9 then\n String.concat \"\" (Array.to_list (Array.map (fun _ -> blocks.(4)) values))\n else\n String.concat \"\"\n (Array.to_list\n (Array.map\n (fun v ->\n let idx = Float.to_int ((v -. lo) /. range *. 7.0) in\n blocks.(max 0 (min 7 idx)))\n values))\n\n(* Q-table *)\n\nlet n_states = n_bins * n_bins * n_bins * n_bins\nlet q = Array.make (n_states * n_actions) 0.0\nlet q_get s a = q.((s * n_actions) + a)\nlet q_set s a v = q.((s * n_actions) + a) <- v\n\n(* Discretize: clip each of the 4 obs dimensions into bins. CartPole obs: [x,\n x_dot, theta, theta_dot] We use generous clip ranges that cover typical\n CartPole trajectories. *)\n\nlet clip_ranges = [| (-2.4, 2.4); (-3.0, 3.0); (-0.21, 0.21); (-3.0, 3.0) |]\n\nlet discretize obs =\n let arr = (Nx.to_array obs : float array) in\n let bin i =\n let lo, hi = clip_ranges.(i) in\n let v = Float.max lo (Float.min hi arr.(i)) in\n let normalized = (v -. lo) /. (hi -. lo) in\n Float.to_int (normalized *. Float.of_int (n_bins - 1))\n |> max 0\n |> min (n_bins - 1)\n in\n let b0 = bin 0 in\n let b1 = bin 1 in\n let b2 = bin 2 in\n let b3 = bin 3 in\n (b0 * n_bins * n_bins * n_bins) + (b1 * n_bins * n_bins) + (b2 * n_bins) + b3\n\nlet best_action s = if q_get s 0 >= q_get s 1 then 0 else 1\n\n(* Training *)\n\nlet () =\n Printf.printf \"Q-Learning on CartPole-v1\\n\";\n Printf.printf \"==========================\\n\\n\";\n Printf.printf \"States: %d bins/dim (%d total), Actions: left/right\\n\" n_bins\n n_states;\n Printf.printf \"alpha = %.2f, gamma = %.2f, episodes = %d\\n\\n\" alpha gamma\n n_episodes;\n\n Nx.Rng.run ~seed:42 @@ fun () ->\n let sample_uniform () =\n let t = Nx.rand Nx.float32 [| 1 |] in\n (Nx.to_array t : float array).(0)\n in\n\n let sample_random_action () =\n let t = Nx.randint Nx.int32 ~high:n_actions [| 1 |] 0 in\n Int32.to_int (Nx.to_array t : Int32.t array).(0)\n in\n\n let env = Fehu_envs.Cartpole.make () in\n\n let n_evals = n_episodes / eval_interval in\n let reward_history = Array.make n_evals 0.0 in\n let eval_idx = ref 0 in\n\n Printf.printf \"Training...\\n\\n\";\n\n for episode = 1 to n_episodes do\n let epsilon =\n epsilon_end\n +. (epsilon_start -. epsilon_end)\n *. exp (-.Float.of_int episode /. epsilon_decay)\n in\n\n let obs, _info = Env.reset env () in\n let state = ref (discretize obs) in\n let done_ = ref false in\n\n while not !done_ do\n let a =\n if sample_uniform () < epsilon then sample_random_action ()\n else best_action !state\n in\n let s = Env.step env (Space.Discrete.of_int a) in\n let next_state = discretize s.observation in\n let done_flag = s.terminated || s.truncated in\n\n let bootstrap =\n if done_flag then 0.0\n else Float.max (q_get next_state 0) (q_get next_state 1)\n in\n let target = s.reward +. (gamma *. bootstrap) in\n let old_q = q_get !state a in\n q_set !state a (old_q +. (alpha *. (target -. old_q)));\n\n state := next_state;\n done_ := done_flag\n done;\n\n if episode mod eval_interval = 0 then begin\n let greedy_policy obs =\n Space.Discrete.of_int (best_action (discretize obs))\n in\n let stats = Eval.run env ~policy:greedy_policy ~n_episodes:20 () in\n Printf.printf\n \" episode %5d eps = %.2f eval: reward = %5.1f +/- %4.1f\\n%!\" episode\n epsilon stats.mean_reward stats.std_reward;\n reward_history.(!eval_idx) <- stats.mean_reward;\n incr eval_idx\n end\n done;\n\n Printf.printf \"\\n reward: %s\\n\" (sparkline reward_history);\n\n (* Final evaluation *)\n Printf.printf \"\\nFinal evaluation (100 episodes):\\n\";\n let greedy_policy obs =\n Space.Discrete.of_int (best_action (discretize obs))\n in\n let stats = Eval.run env ~policy:greedy_policy ~n_episodes:100 () in\n Printf.printf \" mean reward: %5.1f +/- %.1f\\n\" stats.mean_reward\n stats.std_reward;\n Printf.printf \" mean length: %5.1f\\n\" stats.mean_length;\n\n if stats.mean_reward >= 195.0 then\n Printf.printf \"\\nSolved! (mean reward >= 195)\\n\"\n else Printf.printf \"\\nNot solved yet (mean reward < 195).\\n\";\n\n Env.close env\n"
},
{
"path": "packages/fehu/examples/03-reinforce/dune",
"content": "(executable\n (name main)\n (libraries nx rune kaun vega fehu fehu.envs))\n"
},
{
"path": "packages/fehu/examples/03-reinforce/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* REINFORCE on CartPole-v1.\n\n Policy gradient with a small neural network. Collects rollouts, computes\n discounted returns, and updates the policy by maximizing the expected return\n weighted by log-probabilities. *)\n\nopen Fehu\nopen Kaun\n\n(* Hyperparameters *)\n\nlet gamma = 0.99\nlet lr = 1e-3\nlet n_steps = 2048\nlet n_updates = 250\nlet eval_interval = 10\nlet eval_episodes = 20\n\n(* Sparkline *)\n\nlet sparkline values =\n let blocks =\n [|\n \"\\xe2\\x96\\x81\";\n \"\\xe2\\x96\\x82\";\n \"\\xe2\\x96\\x83\";\n \"\\xe2\\x96\\x84\";\n \"\\xe2\\x96\\x85\";\n \"\\xe2\\x96\\x86\";\n \"\\xe2\\x96\\x87\";\n \"\\xe2\\x96\\x88\";\n |]\n in\n let lo = Array.fold_left Float.min Float.infinity values in\n let hi = Array.fold_left Float.max Float.neg_infinity values in\n let range = hi -. lo in\n if range < 1e-9 then\n String.concat \"\" (Array.to_list (Array.map (fun _ -> blocks.(4)) values))\n else\n String.concat \"\"\n (Array.to_list\n (Array.map\n (fun v ->\n let idx = Float.to_int ((v -. lo) /. range *. 7.0) in\n blocks.(max 0 (min 7 idx)))\n values))\n\n(* Network *)\n\nlet network =\n Layer.sequential\n [\n Layer.linear ~in_features:4 ~out_features:64 ();\n Layer.relu ();\n Layer.linear ~in_features:64 ~out_features:2 ();\n ]\n\n(* Forward pass: obs [batch; 4] -> logits [batch; 2] *)\n\nlet forward params net_state obs =\n let vars = Layer.make_vars ~params ~state:net_state ~dtype:Nx.float32 in\n fst (Layer.apply network vars ~training:false obs)\n\n(* Main *)\n\nlet () =\n Printf.printf \"REINFORCE on CartPole-v1\\n\";\n Printf.printf \"=========================\\n\\n\";\n Printf.printf \"Network: Linear(4 -> 64) -> ReLU -> Linear(64 -> 2)\\n\";\n Printf.printf \"Rollout: %d steps/update, gamma = %.2f, lr = %.4f\\n\\n\" n_steps\n gamma lr;\n\n Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make () in\n\n (* Initialize network *)\n let vars = Layer.init network ~dtype:Nx.float32 in\n let params = ref (Layer.params vars) in\n let net_state = Layer.state vars in\n\n Printf.printf \"Parameters: %d\\n\\n\" (Ptree.count_parameters !params);\n\n (* Optimizer *)\n let algo = Vega.adam (Vega.Schedule.constant lr) in\n let opt_state = ref (Optim.init algo !params) in\n\n let policy obs =\n let obs_batch = Nx.reshape [| 1; 4 |] obs in\n let logits = Rune.no_grad (fun () -> forward !params net_state obs_batch) in\n let action_idx = Nx.categorical logits in\n let action = Nx.reshape [||] action_idx in\n let log_probs = Nx.log_softmax logits in\n let action_1 = Nx.reshape [| 1; 1 |] action_idx in\n let log_prob = Nx.take_along_axis ~axis:1 action_1 log_probs in\n let lp = Nx.item [ 0; 0 ] log_prob in\n (action, Some lp, None)\n in\n\n (* Greedy policy for evaluation *)\n let greedy_policy obs =\n let obs_batch = Nx.reshape [| 1; 4 |] obs in\n let logits = Rune.no_grad (fun () -> forward !params net_state obs_batch) in\n let action_idx =\n Nx.argmax logits ~axis:(-1) ~keepdims:false |> Nx.cast Nx.int32\n in\n Nx.reshape [||] action_idx\n in\n\n (* Training loop *)\n Printf.printf \"Training...\\n\\n\";\n\n let n_evals = n_updates / eval_interval in\n let reward_history = Array.make n_evals 0.0 in\n let eval_idx = ref 0 in\n\n for update = 1 to n_updates do\n (* Collect rollout *)\n let traj = Collect.rollout env ~policy ~n_steps in\n let n = Collect.length traj in\n\n (* Compute discounted returns and normalize *)\n let returns =\n Gae.returns ~rewards:traj.rewards ~terminated:traj.terminated\n ~truncated:traj.truncated ~gamma\n in\n let returns = Gae.normalize returns in\n\n (* Stack observations and actions into batch tensors *)\n let obs_batch = Nx.stack (Array.to_list traj.observations) in\n let actions_batch =\n Nx.stack\n (Array.to_list (Array.map (fun a -> Nx.reshape [| 1 |] a) traj.actions))\n in\n let returns_t = Nx.create Nx.float32 [| n |] returns in\n\n (* Policy gradient loss *)\n let loss_fn p =\n let logits = forward p net_state obs_batch in\n let log_probs = Nx.log_softmax logits in\n let action_log_probs =\n Nx.take_along_axis ~axis:1 actions_batch log_probs\n in\n let action_log_probs = Nx.reshape [| n |] action_log_probs in\n let weighted = Nx.mul action_log_probs returns_t in\n Nx.neg (Nx.mean weighted)\n in\n\n let loss, grads = Grad.value_and_grad loss_fn !params in\n let new_params, new_opt_state = Optim.update !opt_state !params grads in\n params := new_params;\n opt_state := new_opt_state;\n\n (* Evaluate periodically *)\n if update mod eval_interval = 0 then begin\n let stats =\n Eval.run env ~policy:greedy_policy ~n_episodes:eval_episodes ()\n in\n Printf.printf\n \" update %3d loss = %6.3f eval: reward = %5.1f +/- %4.1f\\n%!\" update\n (Nx.item [] loss) stats.mean_reward stats.std_reward;\n reward_history.(!eval_idx) <- stats.mean_reward;\n incr eval_idx\n end\n done;\n\n Printf.printf \"\\n reward: %s\\n\" (sparkline reward_history);\n\n (* Final evaluation *)\n Printf.printf \"\\nFinal evaluation (%d episodes):\\n\" 50;\n let stats = Eval.run env ~policy:greedy_policy ~n_episodes:50 () in\n Printf.printf \" mean reward: %5.1f +/- %.1f\\n\" stats.mean_reward\n stats.std_reward;\n Printf.printf \" mean length: %5.1f\\n\" stats.mean_length;\n\n if stats.mean_reward >= 475.0 then\n Printf.printf \"\\nSolved! (mean reward >= 475)\\n\"\n else Printf.printf \"\\nNot solved yet (mean reward < 475).\\n\";\n\n Env.close env\n"
},
{
"path": "packages/fehu/examples/04-dqn/dune",
"content": "(executable\n (name main)\n (libraries nx rune kaun vega fehu fehu.envs))\n"
},
{
"path": "packages/fehu/examples/04-dqn/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* DQN on CartPole-v1.\n\n Deep Q-Network with experience replay and a target network. Epsilon-greedy\n exploration decays linearly. The target network is hard-copied every\n target_update_interval steps. *)\n\nopen Fehu\nopen Kaun\n\n(* Hyperparameters *)\n\nlet buffer_capacity = 50_000\nlet batch_size = 64\nlet gamma = 0.99\nlet lr = 5e-4\nlet epsilon_start = 1.0\nlet epsilon_end = 0.05\nlet epsilon_decay_steps = 10_000\nlet target_update_interval = 250\nlet learning_starts = 1000\nlet n_steps = 50_000\nlet eval_interval = 2000\nlet eval_episodes = 20\n\n(* Sparkline *)\n\nlet sparkline values =\n let blocks =\n [|\n \"\\xe2\\x96\\x81\";\n \"\\xe2\\x96\\x82\";\n \"\\xe2\\x96\\x83\";\n \"\\xe2\\x96\\x84\";\n \"\\xe2\\x96\\x85\";\n \"\\xe2\\x96\\x86\";\n \"\\xe2\\x96\\x87\";\n \"\\xe2\\x96\\x88\";\n |]\n in\n let lo = Array.fold_left Float.min Float.infinity values in\n let hi = Array.fold_left Float.max Float.neg_infinity values in\n let range = hi -. lo in\n if range < 1e-9 then\n String.concat \"\" (Array.to_list (Array.map (fun _ -> blocks.(4)) values))\n else\n String.concat \"\"\n (Array.to_list\n (Array.map\n (fun v ->\n let idx = Float.to_int ((v -. lo) /. range *. 7.0) in\n blocks.(max 0 (min 7 idx)))\n values))\n\n(* Network *)\n\nlet q_network =\n Layer.sequential\n [\n Layer.linear ~in_features:4 ~out_features:128 ();\n Layer.relu ();\n Layer.linear ~in_features:128 ~out_features:128 ();\n Layer.relu ();\n Layer.linear ~in_features:128 ~out_features:2 ();\n ]\n\n(* Forward pass: obs [batch; 4] -> q_values [batch; 2] *)\n\nlet forward params net_state obs =\n let vars = Layer.make_vars ~params ~state:net_state ~dtype:Nx.float32 in\n fst (Layer.apply q_network vars ~training:false obs)\n\n(* Epsilon schedule: linear decay *)\n\nlet epsilon step =\n let t =\n Float.min 1.0 (Float.of_int step /. Float.of_int epsilon_decay_steps)\n in\n epsilon_start +. (t *. (epsilon_end -. epsilon_start))\n\n(* Copy parameters for the target network *)\n\nlet copy_params params = Ptree.map { run = (fun t -> Nx.copy t) } params\n\n(* Main *)\n\nlet () =\n Printf.printf \"DQN on CartPole-v1\\n\";\n Printf.printf \"===================\\n\\n\";\n Printf.printf\n \"Network: Linear(4 -> 128) -> ReLU -> Linear(128 -> 128) -> ReLU -> \\\n Linear(128 -> 2)\\n\";\n Printf.printf \"Buffer: %d, batch: %d, gamma = %.2f, lr = %.4f\\n\"\n buffer_capacity batch_size gamma lr;\n Printf.printf\n \"Epsilon: %.2f -> %.2f over %d steps, target update every %d steps\\n\\n\"\n epsilon_start epsilon_end epsilon_decay_steps target_update_interval;\n\n Nx.Rng.run ~seed:42 @@ fun () ->\n let env = Fehu_envs.Cartpole.make () in\n\n (* Initialize network *)\n let vars = Layer.init q_network ~dtype:Nx.float32 in\n let params = ref (Layer.params vars) in\n let net_state = Layer.state vars in\n let target_params = ref (copy_params !params) in\n\n Printf.printf \"Parameters: %d\\n\\n\" (Ptree.count_parameters !params);\n\n (* Optimizer *)\n let algo = Vega.adam (Vega.Schedule.constant lr) in\n let opt_state = ref (Optim.init algo !params) in\n\n (* Replay buffer *)\n let buffer = Buffer.create ~capacity:buffer_capacity in\n\n let sample_uniform () =\n let t = Nx.rand Nx.float32 [| 1 |] in\n (Nx.to_array t : float array).(0)\n in\n\n (* Epsilon-greedy action selection *)\n let select_action obs eps =\n if sample_uniform () < eps then Space.sample (Env.action_space env)\n else begin\n let obs_batch = Nx.reshape [| 1; 4 |] obs in\n let q_values =\n Rune.no_grad (fun () -> forward !params net_state obs_batch)\n in\n let action_idx =\n Nx.argmax q_values ~axis:(-1) ~keepdims:false |> Nx.cast Nx.int32\n in\n Nx.reshape [||] action_idx\n end\n in\n\n (* Greedy policy for evaluation *)\n let greedy_policy obs =\n let obs_batch = Nx.reshape [| 1; 4 |] obs in\n let q_values =\n Rune.no_grad (fun () -> forward !params net_state obs_batch)\n in\n let action_idx =\n Nx.argmax q_values ~axis:(-1) ~keepdims:false |> Nx.cast Nx.int32\n in\n Nx.reshape [||] action_idx\n in\n\n (* Training step *)\n let train_step () =\n let obs_arr, act_arr, rew_arr, next_obs_arr, term_arr, trunc_arr =\n Buffer.sample_arrays buffer ~batch_size\n in\n let n = Array.length obs_arr in\n\n (* Stack into batch tensors *)\n let obs_batch = Nx.stack (Array.to_list obs_arr) in\n let next_obs_batch = Nx.stack (Array.to_list next_obs_arr) in\n let actions_batch =\n Nx.stack\n (Array.to_list (Array.map (fun a -> Nx.reshape [| 1 |] a) act_arr))\n in\n let rewards_t = Nx.create Nx.float32 [| n |] rew_arr in\n\n (* Done mask: 1.0 if not done, 0.0 if done *)\n let done_mask =\n Array.init n (fun i -> if term_arr.(i) || trunc_arr.(i) then 0.0 else 1.0)\n in\n let done_mask_t = Nx.create Nx.float32 [| n |] done_mask in\n\n (* Compute TD target with target network (no gradient) *)\n let td_target =\n Rune.no_grad (fun () ->\n let target_q = forward !target_params net_state next_obs_batch in\n let max_q = Nx.max target_q ~axes:[ 1 ] ~keepdims:false in\n Nx.add rewards_t\n (Nx.mul (Nx.scalar Nx.float32 gamma) (Nx.mul max_q done_mask_t)))\n in\n let td_target = Rune.detach td_target in\n\n (* Loss: MSE between predicted Q and TD target *)\n let loss_fn p =\n let q_values = forward p net_state obs_batch in\n let q_selected = Nx.take_along_axis ~axis:1 actions_batch q_values in\n let q_selected = Nx.reshape [| n |] q_selected in\n let diff = Nx.sub q_selected td_target in\n Nx.mean (Nx.mul diff diff)\n in\n\n let loss, grads = Grad.value_and_grad loss_fn !params in\n let new_params, new_opt_state = Optim.update !opt_state !params grads in\n params := new_params;\n opt_state := new_opt_state;\n Nx.item [] loss\n in\n\n (* Main training loop *)\n Printf.printf \"Filling buffer (%d steps)...\\n\\n\" learning_starts;\n\n let obs = ref (fst (Env.reset env ())) in\n let last_loss = ref 0.0 in\n\n let n_evals = n_steps / eval_interval in\n let reward_history = Array.make n_evals 0.0 in\n let eval_idx = ref 0 in\n\n Printf.printf \"Training...\\n\\n\";\n\n for step = 1 to n_steps do\n let eps = epsilon step in\n let action = select_action !obs eps in\n let s = Env.step env action in\n\n Buffer.add buffer\n {\n observation = !obs;\n action;\n reward = s.reward;\n next_observation = s.observation;\n terminated = s.terminated;\n truncated = s.truncated;\n };\n\n if s.terminated || s.truncated then obs := fst (Env.reset env ())\n else obs := s.observation;\n\n (* Train *)\n if step >= learning_starts then begin\n last_loss := train_step ();\n\n (* Update target network *)\n if step mod target_update_interval = 0 then\n target_params := copy_params !params\n end;\n\n (* Evaluate periodically *)\n if step mod eval_interval = 0 then begin\n let stats =\n Eval.run env ~policy:greedy_policy ~n_episodes:eval_episodes ()\n in\n Printf.printf\n \" step %5d epsilon = %.2f loss = %6.4f eval: reward = %5.1f +/- \\\n %4.1f\\n\\\n %!\"\n step eps !last_loss stats.mean_reward stats.std_reward;\n reward_history.(!eval_idx) <- stats.mean_reward;\n incr eval_idx;\n (* Eval.run leaves the env in a done state; reset for training *)\n obs := fst (Env.reset env ())\n end\n done;\n\n Printf.printf \"\\n reward: %s\\n\" (sparkline reward_history);\n\n (* Final evaluation *)\n Printf.printf \"\\nFinal evaluation (%d episodes):\\n\" 50;\n let stats = Eval.run env ~policy:greedy_policy ~n_episodes:50 () in\n Printf.printf \" mean reward: %5.1f +/- %.1f\\n\" stats.mean_reward\n stats.std_reward;\n Printf.printf \" mean length: %5.1f\\n\" stats.mean_length;\n\n if stats.mean_reward >= 475.0 then\n Printf.printf \"\\nSolved! (mean reward >= 475)\\n\"\n else Printf.printf \"\\nNot solved yet (mean reward < 475).\\n\";\n\n Env.close env\n"
},
{
"path": "packages/fehu/lib/buffer.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet err_capacity = \"Buffer.create: capacity must be positive\"\nlet err_empty = \"Buffer.sample: buffer is empty\"\nlet err_batch_size = \"Buffer.sample: batch_size must be positive\"\n\ntype ('obs, 'act) transition = {\n observation : 'obs;\n action : 'act;\n reward : float;\n next_observation : 'obs;\n terminated : bool;\n truncated : bool;\n}\n\ntype ('obs, 'act) t = {\n capacity : int;\n mutable size : int;\n mutable pos : int;\n mutable observations : 'obs array;\n mutable actions : 'act array;\n rewards : float array;\n mutable next_observations : 'obs array;\n terminateds : bool array;\n truncateds : bool array;\n}\n\n(* Constructor *)\n\nlet create ~capacity =\n if capacity <= 0 then invalid_arg err_capacity;\n {\n capacity;\n size = 0;\n pos = 0;\n observations = [||];\n actions = [||];\n rewards = Array.make capacity 0.0;\n next_observations = [||];\n terminateds = Array.make capacity false;\n truncateds = Array.make capacity false;\n }\n\n(* Mutating *)\n\nlet ensure_init buf (tr : _ transition) =\n if Array.length buf.observations = 0 then begin\n buf.observations <- Array.make buf.capacity tr.observation;\n buf.actions <- Array.make buf.capacity tr.action;\n buf.next_observations <- Array.make buf.capacity tr.next_observation\n end\n\nlet add buf tr =\n ensure_init buf tr;\n buf.observations.(buf.pos) <- tr.observation;\n buf.actions.(buf.pos) <- tr.action;\n buf.rewards.(buf.pos) <- tr.reward;\n buf.next_observations.(buf.pos) <- tr.next_observation;\n buf.terminateds.(buf.pos) <- tr.terminated;\n buf.truncateds.(buf.pos) <- tr.truncated;\n buf.pos <- (buf.pos + 1) mod buf.capacity;\n if buf.size < buf.capacity then buf.size <- buf.size + 1\n\nlet clear buf =\n buf.size <- 0;\n buf.pos <- 0\n\n(* Sampling *)\n\nlet sample_indices buf ~batch_size =\n if buf.size = 0 then invalid_arg err_empty;\n if batch_size <= 0 then invalid_arg err_batch_size;\n let n = min batch_size buf.size in\n let raw = Nx.randint Nx.int32 ~high:buf.size [| n |] 0 in\n let idx : Int32.t array = Nx.to_array raw in\n (idx, n)\n\nlet sample buf ~batch_size =\n let idx, n = sample_indices buf ~batch_size in\n Array.init n (fun i ->\n let j = Int32.to_int idx.(i) in\n {\n observation = buf.observations.(j);\n action = buf.actions.(j);\n reward = buf.rewards.(j);\n next_observation = buf.next_observations.(j);\n terminated = buf.terminateds.(j);\n truncated = buf.truncateds.(j);\n })\n\nlet sample_arrays buf ~batch_size =\n let idx, n = sample_indices buf ~batch_size in\n let get arr i = arr.(Int32.to_int idx.(i)) in\n let observations = Array.init n (get buf.observations) in\n let actions = Array.init n (get buf.actions) in\n let rewards = Array.init n (get buf.rewards) in\n let next_observations = Array.init n (get buf.next_observations) in\n let terminated = Array.init n (get buf.terminateds) in\n let truncated = Array.init n (get buf.truncateds) in\n (observations, actions, rewards, next_observations, terminated, truncated)\n\n(* Queries *)\n\nlet size buf = buf.size\nlet is_full buf = buf.size = buf.capacity\nlet capacity buf = buf.capacity\n"
},
{
"path": "packages/fehu/lib/buffer.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Replay buffer for off-policy experience storage.\n\n A fixed-capacity circular buffer that stores transitions and supports\n uniform random sampling. Observation and action arrays are lazily\n initialized on the first {!add}. *)\n\n(** {1:types Types} *)\n\ntype ('obs, 'act) transition = {\n observation : 'obs; (** State before the action. *)\n action : 'act; (** Action taken. *)\n reward : float; (** Scalar reward received. *)\n next_observation : 'obs; (** State after the action. *)\n terminated : bool; (** Natural episode ending. *)\n truncated : bool; (** Forced episode ending. *)\n}\n(** The type for transitions. *)\n\ntype ('obs, 'act) t\n(** A replay buffer of transitions. *)\n\n(** {1:constructors Constructors} *)\n\nval create : capacity:int -> ('obs, 'act) t\n(** [create ~capacity] is an empty buffer that holds at most [capacity]\n transitions.\n\n Raises [Invalid_argument] if [capacity <= 0]. *)\n\n(** {1:mutating Mutating} *)\n\nval add : ('obs, 'act) t -> ('obs, 'act) transition -> unit\n(** [add buf tr] appends [tr], overwriting the oldest transition when at\n capacity. *)\n\nval clear : ('obs, 'act) t -> unit\n(** [clear buf] removes all transitions, keeping storage allocated. *)\n\n(** {1:sampling Sampling} *)\n\nval sample : ('obs, 'act) t -> batch_size:int -> ('obs, 'act) transition array\n(** [sample buf ~batch_size] draws [batch_size] transitions uniformly at random\n (with replacement).\n\n Random keys are drawn from the implicit RNG scope.\n\n If [batch_size] exceeds {!size}, samples [min batch_size size] transitions.\n\n Raises [Invalid_argument] if [buf] is empty or [batch_size <= 0]. *)\n\nval sample_arrays :\n ('obs, 'act) t ->\n batch_size:int ->\n 'obs array * 'act array * float array * 'obs array * bool array * bool array\n(** [sample_arrays buf ~batch_size] is like {!sample} but returns\n structure-of-arrays\n [(observations, actions, rewards, next_observations, terminated, truncated)]\n for direct use in training loops. *)\n\n(** {1:queries Queries} *)\n\nval size : ('obs, 'act) t -> int\n(** [size buf] is the number of stored transitions. *)\n\nval is_full : ('obs, 'act) t -> bool\n(** [is_full buf] is [true] iff [size buf = capacity]. *)\n\nval capacity : ('obs, 'act) t -> int\n(** [capacity buf] is the maximum number of transitions [buf] can hold. *)\n"
},
{
"path": "packages/fehu/lib/collect.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet err_concat_empty = \"Collect.concat: empty list\"\n\ntype ('obs, 'act) t = {\n observations : 'obs array;\n actions : 'act array;\n rewards : float array;\n next_observations : 'obs array;\n terminated : bool array;\n truncated : bool array;\n infos : Info.t array;\n log_probs : float array option;\n values : float array option;\n}\n\nlet length t = Array.length t.observations\n\n(* Concatenation *)\n\nlet concat_opt_field ts get =\n if List.for_all (fun t -> Option.is_some (get t)) ts then\n Some (Array.concat (List.map (fun t -> Option.get (get t)) ts))\n else None\n\nlet concat = function\n | [] -> invalid_arg err_concat_empty\n | [ t ] -> t\n | ts ->\n {\n observations = Array.concat (List.map (fun t -> t.observations) ts);\n actions = Array.concat (List.map (fun t -> t.actions) ts);\n rewards = Array.concat (List.map (fun t -> t.rewards) ts);\n next_observations =\n Array.concat (List.map (fun t -> t.next_observations) ts);\n terminated = Array.concat (List.map (fun t -> t.terminated) ts);\n truncated = Array.concat (List.map (fun t -> t.truncated) ts);\n infos = Array.concat (List.map (fun t -> t.infos) ts);\n log_probs = concat_opt_field ts (fun t -> t.log_probs);\n values = concat_opt_field ts (fun t -> t.values);\n }\n\n(* Accumulator for building trajectories *)\n\ntype ('obs, 'act) acc = {\n mutable obs : 'obs list;\n mutable acts : 'act list;\n mutable rews : float list;\n mutable next_obs : 'obs list;\n mutable terms : bool list;\n mutable truncs : bool list;\n mutable infos_acc : Info.t list;\n mutable lps : float list;\n mutable vals : float list;\n mutable count : int;\n}\n\nlet create_acc () =\n {\n obs = [];\n acts = [];\n rews = [];\n next_obs = [];\n terms = [];\n truncs = [];\n infos_acc = [];\n lps = [];\n vals = [];\n count = 0;\n }\n\nlet acc_step acc ~current_obs ~action ~lp_opt ~v_opt (s : _ Env.step) =\n acc.obs <- current_obs :: acc.obs;\n acc.acts <- action :: acc.acts;\n acc.rews <- s.reward :: acc.rews;\n acc.next_obs <- s.observation :: acc.next_obs;\n acc.terms <- s.terminated :: acc.terms;\n acc.truncs <- s.truncated :: acc.truncs;\n acc.infos_acc <- s.info :: acc.infos_acc;\n (match lp_opt with Some lp -> acc.lps <- lp :: acc.lps | None -> ());\n (match v_opt with Some v -> acc.vals <- v :: acc.vals | None -> ());\n acc.count <- acc.count + 1\n\nlet acc_to_trajectory acc =\n let n = acc.count in\n let log_probs =\n if List.length acc.lps = n then Some (Array.of_list (List.rev acc.lps))\n else None\n in\n let values =\n if List.length acc.vals = n then Some (Array.of_list (List.rev acc.vals))\n else None\n in\n {\n observations = Array.of_list (List.rev acc.obs);\n actions = Array.of_list (List.rev acc.acts);\n rewards = Array.of_list (List.rev acc.rews);\n next_observations = Array.of_list (List.rev acc.next_obs);\n terminated = Array.of_list (List.rev acc.terms);\n truncated = Array.of_list (List.rev acc.truncs);\n infos = Array.of_list (List.rev acc.infos_acc);\n log_probs;\n values;\n }\n\n(* Collecting *)\n\nlet rollout env ~policy ~n_steps =\n let acc = create_acc () in\n let obs, _info = Env.reset env () in\n let current_obs = ref obs in\n while acc.count < n_steps do\n let action, lp_opt, v_opt = policy !current_obs in\n let s = Env.step env action in\n acc_step acc ~current_obs:!current_obs ~action ~lp_opt ~v_opt s;\n current_obs := s.observation;\n if s.terminated || s.truncated then begin\n let obs, _info = Env.reset env () in\n current_obs := obs\n end\n done;\n acc_to_trajectory acc\n\nlet episodes env ~policy ~n_episodes ?(max_steps = 1000) () =\n let eps = ref [] in\n for _ = 1 to n_episodes do\n let acc = create_acc () in\n let obs, _info = Env.reset env () in\n let current_obs = ref obs in\n let done_flag = ref false in\n while acc.count < max_steps && not !done_flag do\n let action, lp_opt, v_opt = policy !current_obs in\n let s = Env.step env action in\n acc_step acc ~current_obs:!current_obs ~action ~lp_opt ~v_opt s;\n current_obs := s.observation;\n done_flag := s.terminated || s.truncated\n done;\n eps := acc_to_trajectory acc :: !eps\n done;\n List.rev !eps\n"
},
{
"path": "packages/fehu/lib/collect.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Trajectory collection from environments.\n\n Collects sequential agent-environment interactions into structure-of-arrays\n form for batch processing. Handles automatic resets on episode boundaries\n and records both the current and next observation at each timestep. *)\n\n(** {1:types Types} *)\n\ntype ('obs, 'act) t = {\n observations : 'obs array; (** States before each action. *)\n actions : 'act array; (** Actions taken. *)\n rewards : float array; (** Scalar rewards received. *)\n next_observations : 'obs array; (** States after each action. *)\n terminated : bool array; (** Natural episode endings. *)\n truncated : bool array; (** Forced episode endings. *)\n infos : Info.t array; (** Per-step metadata. *)\n log_probs : float array option; (** Policy log-probabilities. *)\n values : float array option; (** Value estimates. *)\n}\n(** The type for trajectories. All arrays have the same length. Optional fields\n are [None] when the policy does not provide them. *)\n\n(** {1:accessors Accessors} *)\n\nval length : ('obs, 'act) t -> int\n(** [length traj] is the number of transitions in [traj]. *)\n\n(** {1:combining Combining} *)\n\nval concat : ('obs, 'act) t list -> ('obs, 'act) t\n(** [concat trajs] concatenates [trajs] into a single trajectory. Optional\n fields are kept only if present in all inputs.\n\n Raises [Invalid_argument] if [trajs] is empty. *)\n\n(** {1:collecting Collecting} *)\n\nval rollout :\n ('obs, 'act, 'render) Env.t ->\n policy:('obs -> 'act * float option * float option) ->\n n_steps:int ->\n ('obs, 'act) t\n(** [rollout env ~policy ~n_steps] collects [n_steps] transitions.\n\n Resets [env] at the start and automatically on episode boundaries\n (terminated or truncated). The [policy] receives the current observation and\n returns [(action, log_prob_opt, value_opt)]. *)\n\nval episodes :\n ('obs, 'act, 'render) Env.t ->\n policy:('obs -> 'act * float option * float option) ->\n n_episodes:int ->\n ?max_steps:int ->\n unit ->\n ('obs, 'act) t list\n(** [episodes env ~policy ~n_episodes ()] collects complete episodes, one\n trajectory per episode. Each episode runs until termination, truncation, or\n [max_steps] (default [1000]). *)\n"
},
{
"path": "packages/fehu/lib/dune",
"content": "(library\n (name fehu)\n (public_name fehu)\n (libraries nx))\n"
},
{
"path": "packages/fehu/lib/env.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet strf = Printf.sprintf\n\n(* Error messages *)\n\nlet err_closed op = strf \"Env: operation '%s' on a closed environment\" op\n\nlet err_needs_reset op =\n strf \"Env: operation '%s' requires calling reset first\" op\n\nlet err_render_mode mode modes =\n strf \"Env.create: render mode '%s' not in render_modes [%s]\" mode\n (String.concat \"; \" modes)\n\nlet err_obs_reset value =\n strf \"Env.reset: observation outside observation_space (value=%s)\" value\n\nlet err_obs_step value =\n strf \"Env.step: observation outside observation_space (value=%s)\" value\n\nlet err_action value =\n strf \"Env.step: action outside action_space (value=%s)\" value\n\n(* Step result *)\n\ntype 'obs step = {\n observation : 'obs;\n reward : float;\n terminated : bool;\n truncated : bool;\n info : Info.t;\n}\n\nlet step_result ~observation ?(reward = 0.) ?(terminated = false)\n ?(truncated = false) ?(info = Info.empty) () =\n { observation; reward; terminated; truncated; info }\n\n(* Render mode *)\n\ntype render_mode = [ `Human | `Rgb_array | `Ansi | `Svg | `Custom of string ]\n\nlet render_mode_to_string = function\n | `Human -> \"human\"\n | `Rgb_array -> \"rgb_array\"\n | `Ansi -> \"ansi\"\n | `Svg -> \"svg\"\n | `Custom name -> name\n\n(* Shared mutable state *)\n\ntype shared = { mutable closed : bool; mutable needs_reset : bool }\n\n(* Environment *)\n\ntype ('obs, 'act, 'render) t = {\n id : string option;\n observation_space : 'obs Space.t;\n action_space : 'act Space.t;\n render_mode : render_mode option;\n render_modes : string list;\n shared : shared;\n reset_fn : ?options:Info.t -> unit -> 'obs * Info.t;\n step_fn : 'act -> 'obs step;\n render_fn : unit -> 'render option;\n close_fn : unit -> unit;\n}\n\n(* Lifecycle guards *)\n\nlet ensure_open shared op = if shared.closed then invalid_arg (err_closed op)\n\nlet ensure_reset shared op =\n if shared.needs_reset then invalid_arg (err_needs_reset op)\n\n(* Constructor *)\n\nlet create ?id ~observation_space ~action_space ?render_mode\n ?(render_modes = []) ~reset ~step ?render ?close () =\n (match render_mode with\n | None -> ()\n | Some mode ->\n let mode_s = render_mode_to_string mode in\n if not (List.mem mode_s render_modes) then\n invalid_arg (err_render_mode mode_s render_modes));\n let shared = { closed = false; needs_reset = true } in\n let render_fn = Option.value render ~default:(fun () -> None) in\n let close_fn = Option.value close ~default:(fun () -> ()) in\n let rec env =\n {\n id;\n observation_space;\n action_space;\n render_mode;\n render_modes;\n shared;\n reset_fn = (fun ?options () -> reset env ?options ());\n step_fn = (fun action -> step env action);\n render_fn;\n close_fn;\n }\n in\n env\n\n(* Wrap *)\n\nlet wrap ?id ~observation_space ~action_space ?render_mode ~reset ~step ?render\n ?close inner =\n let render_mode =\n match render_mode with Some _ -> render_mode | None -> inner.render_mode\n in\n let render_fn =\n match render with\n | Some f -> fun () -> f inner\n | None -> fun () -> inner.render_fn ()\n in\n let close_fn =\n match close with\n | Some f -> fun () -> f inner\n | None -> fun () -> inner.close_fn ()\n in\n {\n id;\n observation_space;\n action_space;\n render_mode;\n render_modes = inner.render_modes;\n shared = inner.shared;\n reset_fn = (fun ?options () -> reset inner ?options ());\n step_fn = (fun action -> step inner action);\n render_fn;\n close_fn;\n }\n\n(* Accessors *)\n\nlet id env = env.id\nlet observation_space env = env.observation_space\nlet action_space env = env.action_space\nlet render_mode env = env.render_mode\n\n(* Human render helper *)\n\nlet maybe_human_render env =\n match env.render_mode with\n | Some `Human -> ignore (env.render_fn ())\n | _ -> ()\n\n(* Lifecycle — all guards live here *)\n\nlet closed env = env.shared.closed\n\nlet reset env ?options () =\n ensure_open env.shared \"reset\";\n let observation, info = env.reset_fn ?options () in\n if not (Space.contains env.observation_space observation) then\n invalid_arg\n (err_obs_reset\n (Space.pack env.observation_space observation |> Value.to_string));\n env.shared.needs_reset <- false;\n maybe_human_render env;\n (observation, info)\n\nlet step env action =\n ensure_open env.shared \"step\";\n ensure_reset env.shared \"step\";\n if not (Space.contains env.action_space action) then\n invalid_arg\n (err_action (Space.pack env.action_space action |> Value.to_string));\n let result = env.step_fn action in\n if not (Space.contains env.observation_space result.observation) then\n invalid_arg\n (err_obs_step\n (Space.pack env.observation_space result.observation |> Value.to_string));\n if result.terminated || result.truncated then env.shared.needs_reset <- true;\n maybe_human_render env;\n result\n\nlet render env =\n ensure_open env.shared \"render\";\n env.render_fn ()\n\nlet close env =\n if not env.shared.closed then begin\n env.close_fn ();\n env.shared.closed <- true;\n env.shared.needs_reset <- true\n end\n\n(* Wrapper helpers *)\n\nlet err_clip_bounds = \"Env.clip_action: mismatched low/high bounds\"\nlet err_clip_obs_bounds = \"Env.clip_observation: mismatched low/high bounds\"\nlet err_time_limit = \"Env.time_limit: max_episode_steps must be positive\"\n\nlet derive_id env suffix =\n match env.id with None -> None | Some id -> Some (id ^ suffix)\n\nlet clamp_tensor ~low ~high tensor =\n let data = Nx.to_array tensor in\n let clipped = Array.copy data in\n let upper = Array.length clipped - 1 in\n for idx = 0 to upper do\n let lo = low.(idx) in\n let hi = high.(idx) in\n let v = clipped.(idx) in\n if v < lo then clipped.(idx) <- lo else if v > hi then clipped.(idx) <- hi\n done;\n Nx.create Nx.float32 (Nx.shape tensor) clipped\n\n(* Wrappers *)\n\nlet map_observation ~observation_space ~f env =\n wrap\n ?id:(derive_id env \"/ObservationWrapper\")\n ~observation_space ~action_space:env.action_space\n ~reset:(fun inner ?options () ->\n let obs, info = reset inner ?options () in\n f obs info)\n ~step:(fun inner action ->\n let s = step inner action in\n let obs, info = f s.observation s.info in\n { s with observation = obs; info })\n env\n\nlet map_action ~action_space ~f env =\n wrap\n ?id:(derive_id env \"/ActionWrapper\")\n ~observation_space:env.observation_space ~action_space\n ~reset:(fun inner ?options () -> reset inner ?options ())\n ~step:(fun inner action ->\n let s = step inner (f action) in\n {\n observation = s.observation;\n reward = s.reward;\n terminated = s.terminated;\n truncated = s.truncated;\n info = s.info;\n })\n env\n\nlet map_reward ~f env =\n wrap\n ?id:(derive_id env \"/RewardWrapper\")\n ~observation_space:env.observation_space ~action_space:env.action_space\n ~reset:(fun inner ?options () -> reset inner ?options ())\n ~step:(fun inner action ->\n let s = step inner action in\n let reward, info = f ~reward:s.reward ~info:s.info in\n { s with reward; info })\n env\n\n(* Clipping *)\n\nlet clip_action env =\n let low, high = Space.Box.bounds env.action_space in\n let element_count = Array.length low in\n if Array.length high <> element_count then invalid_arg err_clip_bounds;\n let relaxed_low =\n Array.init element_count (fun i ->\n if Float.equal low.(i) high.(i) then low.(i) else Float.neg_infinity)\n in\n let relaxed_high =\n Array.init element_count (fun i ->\n if Float.equal low.(i) high.(i) then high.(i) else Float.infinity)\n in\n let relaxed_space = Space.Box.create ~low:relaxed_low ~high:relaxed_high in\n map_action ~action_space:relaxed_space\n ~f:(fun action -> clamp_tensor ~low ~high action)\n env\n\nlet clip_observation ~low ~high env =\n let inner_low, inner_high = Space.Box.bounds env.observation_space in\n let n = Array.length low in\n if Array.length high <> n then invalid_arg err_clip_obs_bounds;\n if Array.length inner_low <> n then invalid_arg err_clip_obs_bounds;\n let clamp_low = Array.init n (fun i -> Float.max low.(i) inner_low.(i)) in\n let clamp_high = Array.init n (fun i -> Float.min high.(i) inner_high.(i)) in\n let observation_space = Space.Box.create ~low:clamp_low ~high:clamp_high in\n map_observation ~observation_space\n ~f:(fun obs info ->\n (clamp_tensor ~low:clamp_low ~high:clamp_high obs, info))\n env\n\n(* Limits *)\n\nlet time_limit ~max_episode_steps env =\n if max_episode_steps <= 0 then invalid_arg err_time_limit;\n let steps = ref 0 in\n let add_info info elapsed =\n info\n |> Info.set \"time_limit.truncated\" (Info.bool true)\n |> Info.set \"time_limit.elapsed_steps\" (Info.int elapsed)\n in\n wrap\n ?id:(derive_id env \"/TimeLimit\")\n ~observation_space:env.observation_space ~action_space:env.action_space\n ~reset:(fun inner ?options () ->\n steps := 0;\n reset inner ?options ())\n ~step:(fun inner action ->\n incr steps;\n let s = step inner action in\n if s.terminated || s.truncated then begin\n steps := 0;\n s\n end\n else if !steps >= max_episode_steps then begin\n let info = add_info s.info !steps in\n steps := 0;\n { s with truncated = true; info }\n end\n else s)\n env\n"
},
{
"path": "packages/fehu/lib/env.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Reinforcement learning environments.\n\n An environment defines an interactive loop: the agent observes, acts, and\n receives a reward. The environment enforces a lifecycle: {!reset} must be\n called before {!step}, and a terminated or truncated episode requires\n another {!reset}. *)\n\n(** {1:step Step results} *)\n\ntype 'obs step = {\n observation : 'obs; (** The observation after the action. *)\n reward : float; (** Scalar reward for the transition. *)\n terminated : bool; (** [true] when the episode ends naturally. *)\n truncated : bool; (** [true] when the episode is cut short. *)\n info : Info.t; (** Auxiliary metadata. *)\n}\n(** The type for step results. *)\n\nval step_result :\n observation:'obs ->\n ?reward:float ->\n ?terminated:bool ->\n ?truncated:bool ->\n ?info:Info.t ->\n unit ->\n 'obs step\n(** [step_result ~observation ()] constructs a step result. [reward] defaults to\n [0.], [terminated] and [truncated] default to [false], [info] defaults to\n {!Info.empty}. *)\n\n(** {1:render Render modes} *)\n\ntype render_mode = [ `Human | `Rgb_array | `Ansi | `Svg | `Custom of string ]\n(** Rendering modes supported by environments. *)\n\nval render_mode_to_string : render_mode -> string\n(** [render_mode_to_string m] is the string representation of [m]. *)\n\n(** {1:env Environments} *)\n\ntype ('obs, 'act, 'render) t\n(** Environment handle. Use {!create} or {!wrap} to construct. *)\n\nval create :\n ?id:string ->\n observation_space:'obs Space.t ->\n action_space:'act Space.t ->\n ?render_mode:render_mode ->\n ?render_modes:string list ->\n reset:(('obs, 'act, 'render) t -> ?options:Info.t -> unit -> 'obs * Info.t) ->\n step:(('obs, 'act, 'render) t -> 'act -> 'obs step) ->\n ?render:(unit -> 'render option) ->\n ?close:(unit -> unit) ->\n unit ->\n ('obs, 'act, 'render) t\n(** [create ~observation_space ~action_space ~reset ~step ()] makes a new\n environment.\n\n [reset] and [step] receive the environment handle as first argument. Random\n keys for stochastic behavior are drawn from the implicit RNG scope.\n\n [render_modes] lists the supported render mode strings. When [render_mode]\n is provided, it must appear in [render_modes].\n\n Raises [Invalid_argument] if [render_mode] is not in [render_modes]. *)\n\nval wrap :\n ?id:string ->\n observation_space:'obs2 Space.t ->\n action_space:'act2 Space.t ->\n ?render_mode:render_mode ->\n reset:(('obs1, 'act1, 'render) t -> ?options:Info.t -> unit -> 'obs2 * Info.t) ->\n step:(('obs1, 'act1, 'render) t -> 'act2 -> 'obs2 step) ->\n ?render:(('obs1, 'act1, 'render) t -> 'render option) ->\n ?close:(('obs1, 'act1, 'render) t -> unit) ->\n ('obs1, 'act1, 'render) t ->\n ('obs2, 'act2, 'render) t\n(** [wrap ~observation_space ~action_space ~reset ~step inner] builds a new\n environment that wraps [inner]. The wrapper shares [inner]'s lifecycle state\n (RNG, closed flag, reset flag). All guards (closed, needs-reset, space\n validation) are enforced by {!reset}/{!step}, so wrappers get them\n automatically.\n\n The render type is preserved from [inner]. [render_mode] defaults to\n [inner]'s. *)\n\n(** {1:accessors Accessors} *)\n\nval id : ('obs, 'act, 'render) t -> string option\n(** [id env] is the environment's identifier, if any. *)\n\nval observation_space : ('obs, 'act, 'render) t -> 'obs Space.t\n(** [observation_space env] is the space of valid observations. *)\n\nval action_space : ('obs, 'act, 'render) t -> 'act Space.t\n(** [action_space env] is the space of valid actions. *)\n\nval render_mode : ('obs, 'act, 'render) t -> render_mode option\n(** [render_mode env] is the render mode chosen at construction, if any. *)\n\n(** {1:lifecycle Lifecycle} *)\n\nval closed : ('obs, 'act, 'render) t -> bool\n(** [closed env] is [true] iff the environment has been closed. *)\n\nval reset : ('obs, 'act, 'render) t -> ?options:Info.t -> unit -> 'obs * Info.t\n(** [reset env ()] resets the environment to an initial state.\n\n Raises [Invalid_argument] if [env] is closed, or if the reset function\n produces an observation outside {!observation_space}. *)\n\nval step : ('obs, 'act, 'render) t -> 'act -> 'obs step\n(** [step env action] advances the environment by one timestep.\n\n Raises [Invalid_argument] if [env] is closed, if no {!reset} has been called\n since the last terminal step, if [action] is outside {!action_space}, or if\n the step function produces an observation outside {!observation_space}. *)\n\nval render : ('obs, 'act, 'render) t -> 'render option\n(** [render env] produces a visualization of the current state.\n\n Raises [Invalid_argument] if [env] is closed. *)\n\nval close : ('obs, 'act, 'render) t -> unit\n(** [close env] releases resources held by the environment. Subsequent calls are\n no-ops. *)\n\n(** {1:wrappers Wrappers} *)\n\nval map_observation :\n observation_space:'obs2 Space.t ->\n f:('obs1 -> Info.t -> 'obs2 * Info.t) ->\n ('obs1, 'act, 'render) t ->\n ('obs2, 'act, 'render) t\n(** [map_observation ~observation_space ~f env] transforms observations. Every\n observation from {!reset} and {!step} is passed through [f] together with\n the info dictionary. *)\n\nval map_action :\n action_space:'act2 Space.t ->\n f:('act2 -> 'act1) ->\n ('obs, 'act1, 'render) t ->\n ('obs, 'act2, 'render) t\n(** [map_action ~action_space ~f env] transforms actions before passing them to\n the inner environment. *)\n\nval map_reward :\n f:(reward:float -> info:Info.t -> float * Info.t) ->\n ('obs, 'act, 'render) t ->\n ('obs, 'act, 'render) t\n(** [map_reward ~f env] transforms rewards after each step. *)\n\n(** {1:clip Clipping} *)\n\nval clip_action :\n ('obs, Space.Box.element, 'render) t -> ('obs, Space.Box.element, 'render) t\n(** [clip_action env] clamps continuous actions to the bounds of the inner\n environment's {!Space.Box} action space. The wrapper exposes a relaxed space\n that accepts any float values, then clips before forwarding. *)\n\nval clip_observation :\n low:float array ->\n high:float array ->\n (Space.Box.element, 'act, 'render) t ->\n (Space.Box.element, 'act, 'render) t\n(** [clip_observation ~low ~high env] clamps observations to \\[[low]; [high]\\].\n The wrapper's observation space is the intersection of the provided bounds\n and the inner space's bounds.\n\n Raises [Invalid_argument] if [low] and [high] differ in length or do not\n match the inner space's dimensionality. *)\n\n(** {1:limits Limits} *)\n\nval time_limit :\n max_episode_steps:int -> ('obs, 'act, 'render) t -> ('obs, 'act, 'render) t\n(** [time_limit ~max_episode_steps env] enforces a maximum episode length. When\n the limit is reached the step's [truncated] flag is set to [true]. The\n counter resets on {!reset}.\n\n Raises [Invalid_argument] if [max_episode_steps <= 0]. *)\n"
},
{
"path": "packages/fehu/lib/envs/cartpole.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Fehu\n\ntype obs = (float, Nx.float32_elt) Nx.t\ntype act = (int32, Nx.int32_elt) Nx.t\ntype render = string\n\n(* Physics constants matching Gymnasium CartPole-v1 *)\n\nlet gravity = 9.8\nlet masscart = 1.0\nlet masspole = 0.1\nlet total_mass = masscart +. masspole\nlet half_pole_length = 0.5\nlet polemass_length = masspole *. half_pole_length\nlet force_mag = 10.0\nlet tau = 0.02\n\n(* Termination thresholds *)\n\nlet theta_threshold = 12. *. Float.pi /. 180.\nlet x_threshold = 2.4\nlet max_steps = 500\n\n(* Float32-representable large bound for \"unbounded\" dimensions *)\nlet f32_max = 3.4028235e38\n\nlet observation_space =\n Space.Box.create\n ~low:[| -4.8; -.f32_max; -.theta_threshold *. 2.; -.f32_max |]\n ~high:[| 4.8; f32_max; theta_threshold *. 2.; f32_max |]\n\nlet action_space = Space.Discrete.create 2\n\nlet make_obs x x_dot theta theta_dot =\n Nx.create Nx.float32 [| 4 |] [| x; x_dot; theta; theta_dot |]\n\nlet make ?render_mode () =\n let x = ref 0.0 in\n let x_dot = ref 0.0 in\n let theta = ref 0.0 in\n let theta_dot = ref 0.0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n let random_state () =\n let r = Nx.rand Nx.float32 [| 1 |] in\n let v = (Nx.to_array r).(0) in\n (v -. 0.5) *. 0.1\n in\n x := random_state ();\n x_dot := random_state ();\n theta := random_state ();\n theta_dot := random_state ();\n steps := 0;\n (make_obs !x !x_dot !theta !theta_dot, Info.empty)\n in\n let step _env action =\n let force =\n if Space.Discrete.to_int action = 1 then force_mag else -.force_mag\n in\n let costheta = cos !theta in\n let sintheta = sin !theta in\n let temp =\n (force +. (polemass_length *. !theta_dot *. !theta_dot *. sintheta))\n /. total_mass\n in\n let thetaacc =\n ((gravity *. sintheta) -. (costheta *. temp))\n /. (half_pole_length\n *. ((4.0 /. 3.0) -. (masspole *. costheta *. costheta /. total_mass)))\n in\n let xacc =\n temp -. (polemass_length *. thetaacc *. costheta /. total_mass)\n in\n x := !x +. (tau *. !x_dot);\n x_dot := !x_dot +. (tau *. xacc);\n theta := !theta +. (tau *. !theta_dot);\n theta_dot := !theta_dot +. (tau *. thetaacc);\n incr steps;\n let terminated =\n !x < -.x_threshold || !x > x_threshold || !theta < -.theta_threshold\n || !theta > theta_threshold\n in\n let truncated = (not terminated) && !steps >= max_steps in\n let reward = if terminated then 0.0 else 1.0 in\n let info = Info.set \"steps\" (Info.int !steps) Info.empty in\n Env.step_result\n ~observation:(make_obs !x !x_dot !theta !theta_dot)\n ~reward ~terminated ~truncated ~info ()\n in\n let render () =\n Some\n (Printf.sprintf\n \"CartPole: x=%.3f, x_dot=%.3f, theta=%.3f\\xc2\\xb0, theta_dot=%.3f, \\\n steps=%d\"\n !x !x_dot\n (!theta *. 180. /. Float.pi)\n !theta_dot !steps)\n in\n Env.create ?render_mode ~render_modes:[ \"ansi\" ] ~id:\"CartPole-v1\"\n ~observation_space ~action_space ~reset ~step ~render ()\n"
},
{
"path": "packages/fehu/lib/envs/dune",
"content": "(library\n (name fehu_envs)\n (public_name fehu.envs)\n (libraries fehu nx))\n"
},
{
"path": "packages/fehu/lib/envs/fehu_envs.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Random_walk = Random_walk\nmodule Cartpole = Cartpole\nmodule Grid_world = Grid_world\nmodule Mountain_car = Mountain_car\n"
},
{
"path": "packages/fehu/lib/envs/fehu_envs.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Built-in environments for testing and learning.\n\n Four environments covering the standard RL benchmarks: a simple 1D walk, the\n classic cart-pole, a grid navigation problem, and the sparse-reward mountain\n car. All follow the {!Fehu.Env} interface. *)\n\n(** {1:envs Environments} *)\n\nmodule Random_walk : sig\n (** One-dimensional random walk.\n\n The agent moves left or right on a line bounded by \\[[-10]; [10]\\]. Reward\n is [- |position|]. Episodes terminate when the agent reaches a boundary or\n after 200 steps.\n\n {b Observation}: {!Fehu.Space.Box} of shape [[1]] in \\[[-10.0]; [10.0]\\].\n\n {b Actions}: {!Fehu.Space.Discrete} 2 -- 0 = left, 1 = right.\n\n {b Render modes}: [ansi]. *)\n\n type obs = (float, Nx.float32_elt) Nx.t\n type act = (int32, Nx.int32_elt) Nx.t\n type render = string\n\n val make :\n ?render_mode:Fehu.Env.render_mode -> unit -> (obs, act, render) Fehu.Env.t\n (** [make ()] is a random walk environment. *)\nend\n\nmodule Cartpole : sig\n (** Classic cart-pole balancing (CartPole-v1).\n\n A pole is attached to a cart on a frictionless track. The agent pushes the\n cart left or right to keep the pole upright. Reward is [+1.0] per step\n while the pole stays up. The episode terminates when the pole exceeds\n +/-12 degrees or the cart leaves +/-2.4, and truncates at 500 steps.\n\n {b Observation}: {!Fehu.Space.Box} of shape [[4]] -- [x], [x_dot],\n [theta], [theta_dot].\n\n {b Actions}: {!Fehu.Space.Discrete} 2 -- 0 = push left, 1 = push right.\n\n {b Render modes}: [ansi]. *)\n\n type obs = (float, Nx.float32_elt) Nx.t\n type act = (int32, Nx.int32_elt) Nx.t\n type render = string\n\n val make :\n ?render_mode:Fehu.Env.render_mode -> unit -> (obs, act, render) Fehu.Env.t\n (** [make ()] is a cart-pole environment. *)\nend\n\nmodule Grid_world : sig\n (** 5x5 grid navigation with obstacle.\n\n The agent starts at [(0, 0)] and must reach the goal at [(4, 4)]. An\n obstacle at [(2, 2)] blocks movement. Reward is [+10.0] on reaching the\n goal, [-1.0] otherwise. Truncates at 200 steps.\n\n {b Observation}: {!Fehu.Space.Multi_discrete} [[5; 5]] -- [(row, col)].\n\n {b Actions}: {!Fehu.Space.Discrete} 4 -- 0 = up, 1 = down, 2 = left, 3 =\n right.\n\n {b Render modes}: [ansi], [rgb_array]. *)\n\n type obs = (int32, Nx.int32_elt) Nx.t\n type act = (int32, Nx.int32_elt) Nx.t\n type render = Text of string | Image of Fehu.Render.image\n\n val make :\n ?render_mode:Fehu.Env.render_mode -> unit -> (obs, act, render) Fehu.Env.t\n (** [make ()] is a grid world environment. *)\nend\n\nmodule Mountain_car : sig\n (** Mountain car with sparse reward (MountainCar-v0).\n\n A car sits in a valley between two hills. The engine is too weak to climb\n the right hill directly; the agent must build momentum by rocking back and\n forth. Reward is [-1.0] per step. The episode terminates when the car\n reaches position >= 0.5 with non-negative velocity, and truncates at 200\n steps.\n\n {b Observation}: {!Fehu.Space.Box} of shape [[2]] -- [position],\n [velocity].\n\n {b Actions}: {!Fehu.Space.Discrete} 3 -- 0 = push left, 1 = coast, 2 =\n push right.\n\n {b Render modes}: [ansi]. *)\n\n type obs = (float, Nx.float32_elt) Nx.t\n type act = (int32, Nx.int32_elt) Nx.t\n type render = string\n\n val make :\n ?render_mode:Fehu.Env.render_mode -> unit -> (obs, act, render) Fehu.Env.t\n (** [make ()] is a mountain car environment. *)\nend\n"
},
{
"path": "packages/fehu/lib/envs/grid_world.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Fehu\n\ntype obs = (int32, Nx.int32_elt) Nx.t\ntype act = (int32, Nx.int32_elt) Nx.t\ntype render = Text of string | Image of Render.image\n\nlet grid_size = 5\nlet max_steps = 200\nlet observation_space = Space.Multi_discrete.create [| grid_size; grid_size |]\nlet action_space = Space.Discrete.create 4\nlet is_goal row col = row = grid_size - 1 && col = grid_size - 1\nlet is_obstacle row col = row = 2 && col = 2\n\nlet is_valid row col =\n row >= 0 && row < grid_size && col >= 0 && col < grid_size\n && not (is_obstacle row col)\n\nlet make_obs row col =\n Nx.create Nx.int32 [| 2 |] [| Int32.of_int row; Int32.of_int col |]\n\n(* ANSI rendering *)\n\nlet render_text row col =\n let buffer = Bytes.make (grid_size * grid_size) '.' in\n Bytes.set buffer ((row * grid_size) + col) 'A';\n Bytes.set buffer (((grid_size - 1) * grid_size) + (grid_size - 1)) 'G';\n Bytes.set buffer ((2 * grid_size) + 2) '#';\n let rows =\n List.init grid_size (fun r ->\n Bytes.sub_string buffer (r * grid_size) grid_size)\n in\n Format.asprintf \"Position: (%d, %d)@.%a\" row col\n (Format.pp_print_list\n ~pp_sep:(fun fmt () -> Format.fprintf fmt \"@.\")\n Format.pp_print_string)\n rows\n\n(* RGB rendering *)\n\nlet cell_size = 32\nlet frame_width = grid_size * cell_size\nlet frame_height = grid_size * cell_size\n\nlet fill_rect data ~x0 ~y0 ~w ~h ~r ~g ~b =\n for dy = 0 to h - 1 do\n let row_offset = (y0 + dy) * frame_width * 3 in\n for dx = 0 to w - 1 do\n let base = row_offset + ((x0 + dx) * 3) in\n Bigarray.Array1.unsafe_set data base r;\n Bigarray.Array1.unsafe_set data (base + 1) g;\n Bigarray.Array1.unsafe_set data (base + 2) b\n done\n done\n\nlet render_image row col =\n let len = frame_width * frame_height * 3 in\n let data =\n Bigarray.Array1.create Bigarray.int8_unsigned Bigarray.c_layout len\n in\n fill_rect data ~x0:0 ~y0:0 ~w:frame_width ~h:frame_height ~r:30 ~g:33 ~b:36;\n for gr = 0 to grid_size - 1 do\n for gc = 0 to grid_size - 1 do\n let x0 = gc * cell_size in\n let y0 = gr * cell_size in\n fill_rect data ~x0 ~y0 ~w:cell_size ~h:cell_size ~r:44 ~g:48 ~b:52;\n fill_rect data ~x0:(x0 + 1) ~y0:(y0 + 1) ~w:(cell_size - 2)\n ~h:(cell_size - 2) ~r:54 ~g:60 ~b:65\n done\n done;\n let draw_cell cr cc ~r ~g ~b =\n fill_rect data\n ~x0:((cc * cell_size) + 2)\n ~y0:((cr * cell_size) + 2)\n ~w:(cell_size - 4) ~h:(cell_size - 4) ~r ~g ~b\n in\n draw_cell row col ~r:78 ~g:162 ~b:196;\n draw_cell (grid_size - 1) (grid_size - 1) ~r:76 ~g:175 ~b:80;\n draw_cell 2 2 ~r:200 ~g:80 ~b:80;\n Render.image ~width:frame_width ~height:frame_height data\n\nlet make ?render_mode () =\n let row = ref 0 in\n let col = ref 0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n row := 0;\n col := 0;\n steps := 0;\n (make_obs 0 0, Info.empty)\n in\n let step _env action =\n let r, c = (!row, !col) in\n let nr, nc =\n match Space.Discrete.to_int action with\n | 0 -> (r - 1, c)\n | 1 -> (r + 1, c)\n | 2 -> (r, c - 1)\n | 3 -> (r, c + 1)\n | _ -> (r, c)\n in\n let nr, nc = if is_valid nr nc then (nr, nc) else (r, c) in\n row := nr;\n col := nc;\n incr steps;\n let terminated = is_goal nr nc in\n let truncated = (not terminated) && !steps >= max_steps in\n let reward = if terminated then 10.0 else -1.0 in\n let info = Info.set \"steps\" (Info.int !steps) Info.empty in\n Env.step_result ~observation:(make_obs nr nc) ~reward ~terminated ~truncated\n ~info ()\n in\n let render_mode_val = render_mode in\n let render () =\n match render_mode_val with\n | Some `Rgb_array -> Some (Image (render_image !row !col))\n | _ -> Some (Text (render_text !row !col))\n in\n Env.create ?render_mode ~render_modes:[ \"ansi\"; \"rgb_array\" ]\n ~id:\"GridWorld-v0\" ~observation_space ~action_space ~reset ~step ~render ()\n"
},
{
"path": "packages/fehu/lib/envs/mountain_car.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Fehu\n\ntype obs = (float, Nx.float32_elt) Nx.t\ntype act = (int32, Nx.int32_elt) Nx.t\ntype render = string\n\n(* Physics constants matching Gymnasium MountainCar-v0 *)\n\nlet min_position = -1.2\nlet max_position = 0.6\nlet max_speed = 0.07\nlet goal_position = 0.5\nlet goal_velocity = 0.0\nlet force = 0.001\nlet gravity = 0.0025\nlet max_steps = 200\n\nlet observation_space =\n Space.Box.create\n ~low:[| min_position; -.max_speed |]\n ~high:[| max_position; max_speed |]\n\nlet action_space = Space.Discrete.create 3\n\nlet make_obs position velocity =\n Nx.create Nx.float32 [| 2 |] [| position; velocity |]\n\nlet make ?render_mode () =\n let position = ref 0.0 in\n let velocity = ref 0.0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n let r = Nx.rand Nx.float32 [| 1 |] in\n let v = (Nx.to_array r).(0) in\n position := -0.6 +. (v *. 0.2);\n velocity := 0.0;\n steps := 0;\n (make_obs !position !velocity, Info.empty)\n in\n let step _env action =\n let force_direction = float_of_int (Space.Discrete.to_int action - 1) in\n let vel =\n !velocity +. (force_direction *. force)\n -. (gravity *. cos (3.0 *. !position))\n in\n let vel = Float.max (-.max_speed) (Float.min vel max_speed) in\n let pos = !position +. vel in\n let pos = Float.max min_position (Float.min pos max_position) in\n let vel = if pos = min_position && vel < 0.0 then 0.0 else vel in\n position := pos;\n velocity := vel;\n incr steps;\n let terminated = pos >= goal_position && vel >= goal_velocity in\n let truncated = (not terminated) && !steps >= max_steps in\n let reward = -1.0 in\n let info = Info.set \"steps\" (Info.int !steps) Info.empty in\n Env.step_result ~observation:(make_obs pos vel) ~reward ~terminated\n ~truncated ~info ()\n in\n let render () =\n let normalized_pos =\n (!position -. min_position) /. (max_position -. min_position)\n in\n let car_pos = int_of_float (normalized_pos *. 40.0) in\n let goal_pos =\n int_of_float\n ((goal_position -. min_position)\n /. (max_position -. min_position)\n *. 40.0)\n in\n let track = Bytes.make 41 '-' in\n Bytes.set track goal_pos 'G';\n Bytes.set track (max 0 (min 40 car_pos)) 'C';\n Some\n (Printf.sprintf \"MountainCar: [%s] pos=%.3f, vel=%.3f, steps=%d\"\n (Bytes.to_string track) !position !velocity !steps)\n in\n Env.create ?render_mode ~render_modes:[ \"ansi\" ] ~id:\"MountainCar-v0\"\n ~observation_space ~action_space ~reset ~step ~render ()\n"
},
{
"path": "packages/fehu/lib/envs/random_walk.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Fehu\n\ntype obs = (float, Nx.float32_elt) Nx.t\ntype act = (int32, Nx.int32_elt) Nx.t\ntype render = string\n\nlet step_size = 1.0\nlet max_position = 10.0\nlet max_steps = 200\n\nlet observation_space =\n Space.Box.create ~low:[| -.max_position |] ~high:[| max_position |]\n\nlet action_space = Space.Discrete.create 2\nlet make_obs position = Nx.create Nx.float32 [| 1 |] [| position |]\n\nlet render_ansi position =\n let offset = int_of_float (position +. max_position) in\n let offset = max 0 (min 20 offset) in\n let buffer = Bytes.make 21 '.' in\n Bytes.set buffer offset 'o';\n Printf.sprintf \"Position: %+.2f\\n|%s|\" position (Bytes.to_string buffer)\n\nlet make ?render_mode () =\n let position = ref 0.0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n position := 0.0;\n steps := 0;\n (make_obs 0.0, Info.empty)\n in\n let step _env action =\n let direction =\n if Space.Discrete.to_int action = 0 then -.step_size else step_size\n in\n let updated = !position +. direction in\n let clamped = Float.min max_position (Float.max (-.max_position) updated) in\n position := clamped;\n incr steps;\n let terminated = Float.abs clamped >= max_position in\n let truncated = (not terminated) && !steps >= max_steps in\n let reward = -.Float.abs clamped in\n let info = Info.set \"steps\" (Info.int !steps) Info.empty in\n Env.step_result ~observation:(make_obs clamped) ~reward ~terminated\n ~truncated ~info ()\n in\n let render () = Some (render_ansi !position) in\n Env.create ?render_mode ~render_modes:[ \"ansi\" ] ~id:\"RandomWalk-v0\"\n ~observation_space ~action_space ~reset ~step ~render ()\n"
},
{
"path": "packages/fehu/lib/eval.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype stats = {\n mean_reward : float;\n std_reward : float;\n mean_length : float;\n n_episodes : int;\n}\n\nlet run env ~policy ?(n_episodes = 10) ?(max_steps = 1000) () =\n let ep_rewards = Array.make n_episodes 0.0 in\n let ep_lengths = Array.make n_episodes 0.0 in\n for ep = 0 to n_episodes - 1 do\n let obs, _info = Env.reset env () in\n let current_obs = ref obs in\n let total_reward = ref 0.0 in\n let steps = ref 0 in\n let done_flag = ref false in\n while !steps < max_steps && not !done_flag do\n let action = policy !current_obs in\n let s = Env.step env action in\n total_reward := !total_reward +. s.reward;\n steps := !steps + 1;\n current_obs := s.observation;\n done_flag := s.terminated || s.truncated\n done;\n ep_rewards.(ep) <- !total_reward;\n ep_lengths.(ep) <- Float.of_int !steps\n done;\n let n = Float.of_int n_episodes in\n let mean_reward = ref 0.0 in\n let mean_length = ref 0.0 in\n for i = 0 to n_episodes - 1 do\n mean_reward := !mean_reward +. ep_rewards.(i);\n mean_length := !mean_length +. ep_lengths.(i)\n done;\n mean_reward := !mean_reward /. n;\n mean_length := !mean_length /. n;\n let var_sum = ref 0.0 in\n for i = 0 to n_episodes - 1 do\n let d = ep_rewards.(i) -. !mean_reward in\n var_sum := !var_sum +. (d *. d)\n done;\n let std_reward = sqrt (!var_sum /. n) in\n {\n mean_reward = !mean_reward;\n std_reward;\n mean_length = !mean_length;\n n_episodes;\n }\n"
},
{
"path": "packages/fehu/lib/eval.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Policy evaluation.\n\n Runs a deterministic or stochastic policy over multiple episodes and reports\n summary statistics. *)\n\n(** {1:types Types} *)\n\ntype stats = {\n mean_reward : float; (** Mean total reward across episodes. *)\n std_reward : float; (** Standard deviation of total rewards. *)\n mean_length : float; (** Mean episode length in steps. *)\n n_episodes : int; (** Number of episodes evaluated. *)\n}\n(** The type for evaluation statistics. *)\n\n(** {1:running Running} *)\n\nval run :\n ('obs, 'act, 'render) Env.t ->\n policy:('obs -> 'act) ->\n ?n_episodes:int ->\n ?max_steps:int ->\n unit ->\n stats\n(** [run env ~policy ()] evaluates [policy] over [n_episodes] (default [10])\n episodes of at most [max_steps] (default [1000]) steps each. The environment\n is reset between episodes. *)\n"
},
{
"path": "packages/fehu/lib/fehu.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Value = Value\nmodule Info = Info\nmodule Space = Space\nmodule Env = Env\nmodule Vec_env = Vec_env\nmodule Collect = Collect\nmodule Buffer = Buffer\nmodule Gae = Gae\nmodule Eval = Eval\nmodule Render = Render\n"
},
{
"path": "packages/fehu/lib/fehu.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** {0 Fehu} Reinforcement learning environments and utilities.\n\n {1 Core}\n {!modules: Value Info Space Env}\n\n {1 Collection and training}\n {!modules: Collect Buffer Gae Eval}\n\n {1 Composition}\n {!modules: Vec_env Render} *)\n\nmodule Value = Value\nmodule Info = Info\nmodule Space = Space\nmodule Env = Env\nmodule Vec_env = Vec_env\nmodule Collect = Collect\nmodule Buffer = Buffer\nmodule Gae = Gae\nmodule Eval = Eval\nmodule Render = Render\n"
},
{
"path": "packages/fehu/lib/gae.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet err_lengths = \"Gae: all arrays must have the same length\"\n\nlet err_returns_lengths =\n \"Gae.returns: rewards, terminated, and truncated must have the same length\"\n\nlet err_cfv_lengths =\n \"Gae.compute_from_values: all arrays must have the same length\"\n\nlet compute ~rewards ~values ~terminated ~truncated ~next_values ~gamma ~lambda\n =\n let n = Array.length rewards in\n if\n n <> Array.length values\n || n <> Array.length terminated\n || n <> Array.length truncated\n || n <> Array.length next_values\n then invalid_arg err_lengths;\n let advantages = Array.make n 0.0 in\n let returns = Array.make n 0.0 in\n let last_gae = ref 0.0 in\n for t = n - 1 downto 0 do\n let next_val, continuation =\n if terminated.(t) then (0.0, 0.0)\n else if truncated.(t) then (next_values.(t), 0.0)\n else begin\n let v = if t = n - 1 then next_values.(t) else values.(t + 1) in\n (v, 1.0)\n end\n in\n let delta = rewards.(t) +. (gamma *. next_val) -. values.(t) in\n last_gae := delta +. (gamma *. lambda *. continuation *. !last_gae);\n advantages.(t) <- !last_gae;\n returns.(t) <- !last_gae +. values.(t)\n done;\n (advantages, returns)\n\nlet compute_from_values ~rewards ~values ~terminated ~truncated ~last_value\n ~gamma ~lambda =\n let n = Array.length rewards in\n if\n n <> Array.length values\n || n <> Array.length terminated\n || n <> Array.length truncated\n then invalid_arg err_cfv_lengths;\n let next_values =\n Array.init n (fun t -> if t = n - 1 then last_value else values.(t + 1))\n in\n compute ~rewards ~values ~terminated ~truncated ~next_values ~gamma ~lambda\n\nlet returns ~rewards ~terminated ~truncated ~gamma =\n let n = Array.length rewards in\n if n <> Array.length terminated || n <> Array.length truncated then\n invalid_arg err_returns_lengths;\n let ret = Array.make n 0.0 in\n let acc = ref 0.0 in\n for t = n - 1 downto 0 do\n let cont = if terminated.(t) || truncated.(t) then 0.0 else 1.0 in\n acc := rewards.(t) +. (gamma *. cont *. !acc);\n ret.(t) <- !acc\n done;\n ret\n\nlet normalize ?(eps = 1e-8) arr =\n let n = Array.length arr in\n if n = 0 then arr\n else begin\n let mean = ref 0.0 in\n let m2 = ref 0.0 in\n for i = 0 to n - 1 do\n let k = Float.of_int (i + 1) in\n let x = arr.(i) in\n let delta = x -. !mean in\n mean := !mean +. (delta /. k);\n let delta2 = x -. !mean in\n m2 := !m2 +. (delta *. delta2)\n done;\n let std = sqrt (!m2 /. Float.of_int n) +. eps in\n let mu = !mean in\n Array.init n (fun i -> (arr.(i) -. mu) /. std)\n end\n"
},
{
"path": "packages/fehu/lib/gae.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Generalized Advantage Estimation.\n\n Correctly handles the distinction between terminated and truncated episodes.\n On termination, the bootstrap value is zero. On truncation, the bootstrap\n value comes from [next_values]. *)\n\n(** {1:gae GAE} *)\n\nval compute :\n rewards:float array ->\n values:float array ->\n terminated:bool array ->\n truncated:bool array ->\n next_values:float array ->\n gamma:float ->\n lambda:float ->\n float array * float array\n(** [compute ~rewards ~values ~terminated ~truncated ~next_values\n ~gamma ~lambda] is [(advantages, returns)].\n\n [next_values.(t)] is V(s_{{t+1}}). When [terminated.(t)] is\n [true], the bootstrap value is zero and the GAE trace resets.\n When [truncated.(t)] is [true], the bootstrap value is\n [next_values.(t)] and the trace resets for the new episode.\n Otherwise, continuation uses the next step's value.\n\n Raises [Invalid_argument] if array lengths differ. *)\n\nval compute_from_values :\n rewards:float array ->\n values:float array ->\n terminated:bool array ->\n truncated:bool array ->\n last_value:float ->\n gamma:float ->\n lambda:float ->\n float array * float array\n(** [compute_from_values ~rewards ~values ~terminated ~truncated ~last_value\n ~gamma ~lambda] is [(advantages, returns)].\n\n Convenience wrapper around {!compute} that builds [next_values] from\n [values] and [last_value]: [next_values.(t) = values.(t+1)] for [t < n-1],\n and [next_values.(n-1) = last_value].\n\n Raises [Invalid_argument] if array lengths differ. *)\n\n(** {1:returns Monte Carlo returns} *)\n\nval returns :\n rewards:float array ->\n terminated:bool array ->\n truncated:bool array ->\n gamma:float ->\n float array\n(** [returns ~rewards ~terminated ~truncated ~gamma] computes discounted\n cumulative returns. The accumulation resets at terminal or truncated states.\n*)\n\n(** {1:normalize Normalization} *)\n\nval normalize : ?eps:float -> float array -> float array\n(** [normalize arr] is a copy of [arr] with zero mean and unit variance. [eps]\n (default [1e-8]) prevents division by zero. *)\n"
},
{
"path": "packages/fehu/lib/info.ml",
"content": "module String_map = Map.Make (String)\n\ntype t = Value.t String_map.t\n\nlet empty = String_map.empty\nlet is_empty = String_map.is_empty\nlet set key value info = String_map.add key value info\nlet find key info = String_map.find_opt key info\n\nlet find_exn key info =\n match String_map.find_opt key info with\n | Some v -> v\n | None -> invalid_arg (Printf.sprintf \"Info.find_exn: key %S not present\" key)\n\nlet remove key info = String_map.remove key info\nlet merge a b = String_map.union (fun _key _left right -> Some right) a b\nlet to_list info = String_map.bindings info\n\nlet of_list kvs =\n List.fold_left (fun acc (k, v) -> String_map.add k v acc) String_map.empty kvs\n\nlet to_value info = Value.Dict (String_map.bindings info)\n\nlet pp ppf t =\n let bindings = String_map.bindings t in\n Format.fprintf ppf \"{\";\n List.iteri\n (fun i (k, v) ->\n if i > 0 then Format.fprintf ppf \"; \";\n Format.fprintf ppf \"%s: %a\" k Value.pp v)\n bindings;\n Format.fprintf ppf \"}\"\n\n(* Convenience constructors *)\n\nlet null = Value.Null\nlet bool b = Value.Bool b\nlet int i = Value.Int i\nlet float f = Value.Float f\nlet string s = Value.String s\nlet int_array arr = Value.Int_array (Array.copy arr)\nlet float_array arr = Value.Float_array (Array.copy arr)\nlet bool_array arr = Value.Bool_array (Array.copy arr)\n"
},
{
"path": "packages/fehu/lib/info.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Step metadata dictionaries.\n\n Info dictionaries carry auxiliary data returned by {!Env.reset} and\n {!Env.step}. Keys are strings and values are {!Value.t}. *)\n\n(** {1:types Types} *)\n\ntype t\n(** The type for info dictionaries. *)\n\n(** {1:constructors Constructors} *)\n\nval empty : t\n(** [empty] is the empty dictionary. *)\n\nval of_list : (string * Value.t) list -> t\n(** [of_list kvs] is a dictionary from the given key-value pairs. *)\n\n(** {1:predicates Predicates} *)\n\nval is_empty : t -> bool\n(** [is_empty t] is [true] iff [t] has no bindings. *)\n\n(** {1:ops Operations} *)\n\nval set : string -> Value.t -> t -> t\n(** [set k v t] is [t] with [k] bound to [v]. *)\n\nval find : string -> t -> Value.t option\n(** [find k t] is the value bound to [k] in [t], if any. *)\n\nval find_exn : string -> t -> Value.t\n(** [find_exn k t] is the value bound to [k] in [t].\n\n Raises [Invalid_argument] if [k] is not present. *)\n\nval remove : string -> t -> t\n(** [remove k t] is [t] without the binding for [k]. *)\n\nval merge : t -> t -> t\n(** [merge a b] is the union of [a] and [b]. When both have a binding for the\n same key, the value from [b] wins. *)\n\n(** {1:converting Converting} *)\n\nval to_list : t -> (string * Value.t) list\n(** [to_list t] is the bindings of [t] in key order. *)\n\nval to_value : t -> Value.t\n(** [to_value t] is [t] as a {!Value.Dict}. *)\n\n(** {1:fmt Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp] formats an info dictionary for debugging. *)\n\n(** {1:convenience Convenience value constructors} *)\n\nval null : Value.t\n(** [null] is {!Value.Null}. *)\n\nval bool : bool -> Value.t\n(** [bool b] is [Value.Bool b]. *)\n\nval int : int -> Value.t\n(** [int i] is [Value.Int i]. *)\n\nval float : float -> Value.t\n(** [float f] is [Value.Float f]. *)\n\nval string : string -> Value.t\n(** [string s] is [Value.String s]. *)\n\nval int_array : int array -> Value.t\n(** [int_array arr] is [Value.Int_array (Array.copy arr)]. *)\n\nval float_array : float array -> Value.t\n(** [float_array arr] is [Value.Float_array (Array.copy arr)]. *)\n\nval bool_array : bool array -> Value.t\n(** [bool_array arr] is [Value.Bool_array (Array.copy arr)]. *)\n"
},
{
"path": "packages/fehu/lib/render.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Pixel = struct\n type format = Rgb | Rgba | Gray\n\n let channels = function Rgb -> 3 | Rgba -> 4 | Gray -> 1\nend\n\ntype image = {\n width : int;\n height : int;\n pixel_format : Pixel.format;\n data : (int, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t;\n}\n\nlet err_data_length ~expected ~got =\n Printf.sprintf\n \"Render.image: data length %d does not match width * height * channels = %d\"\n got expected\n\nlet image ~width ~height ?(pixel_format = Pixel.Rgb) data =\n let expected = width * height * Pixel.channels pixel_format in\n let got = Bigarray.Array1.dim data in\n if got <> expected then invalid_arg (err_data_length ~expected ~got);\n { width; height; pixel_format; data }\n\nlet rollout env ~policy ~steps ~sink () =\n let obs, _info = Env.reset env () in\n let current_obs = ref obs in\n for _ = 1 to steps do\n let action = policy !current_obs in\n let step = Env.step env action in\n (match Env.render env with Some frame -> sink frame | None -> ());\n current_obs := step.Env.observation;\n if step.Env.terminated || step.Env.truncated then begin\n let obs, _info = Env.reset env () in\n current_obs := obs\n end\n done\n\nlet derive_id env suffix =\n match Env.id env with None -> None | Some id -> Some (id ^ suffix)\n\nlet on_render ~sink env =\n let maybe_record inner =\n match Env.render inner with Some frame -> sink frame | None -> ()\n in\n Env.wrap\n ?id:(derive_id env \"/OnRender\")\n ~observation_space:(Env.observation_space env)\n ~action_space:(Env.action_space env)\n ~reset:(fun inner ?options () ->\n let result = Env.reset inner ?options () in\n maybe_record inner;\n result)\n ~step:(fun inner action ->\n let s = Env.step inner action in\n maybe_record inner;\n s)\n ~render:(fun inner -> Env.render inner)\n env\n"
},
{
"path": "packages/fehu/lib/render.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Visualization primitives.\n\n {!image} is the standard frame type for rgb-rendered environments.\n {!rollout} runs a policy and feeds rendered frames to a user-provided sink.\n*)\n\n(** {1:pixel Pixel formats} *)\n\nmodule Pixel : sig\n (** The type for pixel formats. *)\n type format =\n | Rgb (** 3 channels. *)\n | Rgba (** 4 channels. *)\n | Gray (** 1 channel. *)\n\n val channels : format -> int\n (** [channels fmt] is the number of channels for [fmt]. *)\nend\n\n(** {1:image Images} *)\n\ntype image = {\n width : int; (** Width in pixels. *)\n height : int; (** Height in pixels. *)\n pixel_format : Pixel.format; (** Pixel layout. *)\n data : (int, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t;\n (** Raw pixel data of length [width * height * channels]. *)\n}\n(** The type for rendered frames. *)\n\nval image :\n width:int ->\n height:int ->\n ?pixel_format:Pixel.format ->\n (int, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t ->\n image\n(** [image ~width ~height data] constructs a frame. [pixel_format] defaults to\n [Rgb].\n\n Raises [Invalid_argument] if [Bigarray.Array1.dim data] does not equal\n [width * height * channels]. *)\n\n(** {1:rollout Rollout} *)\n\nval rollout :\n ('obs, 'act, image) Env.t ->\n policy:('obs -> 'act) ->\n steps:int ->\n sink:(image -> unit) ->\n unit ->\n unit\n(** [rollout env ~policy ~steps ~sink ()] runs [policy] in [env] for up to\n [steps] steps. Each rendered frame is passed to [sink]. The environment is\n reset at the start and on episode boundaries. *)\n\n(** {1:recording Recording} *)\n\nval on_render :\n sink:(image -> unit) -> ('obs, 'act, image) Env.t -> ('obs, 'act, image) Env.t\n(** [on_render ~sink env] wraps [env] so that every rendered frame after\n {!Env.reset} and {!Env.step} is passed to [sink]. The wrapper is\n transparent: observations, actions, rewards, and termination signals pass\n through unchanged. *)\n"
},
{
"path": "packages/fehu/lib/space.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Error messages *)\n\nlet err_discrete_n = \"Space.Discrete.create: n must be strictly positive\"\nlet err_discrete_not = \"Space.Discrete: not a discrete space\"\nlet err_box_empty = \"Space.Box.create: low cannot be empty\"\nlet err_box_shape = \"Space.Box.create: low and high must have identical lengths\"\nlet err_box_not = \"Space.Box: not a box space\"\nlet err_mb_n = \"Space.Multi_binary.create: n must be strictly positive\"\nlet err_md_empty = \"Space.Multi_discrete.create: nvec must not be empty\"\nlet err_seq_min = \"Space.Sequence.create: min_length must be non-negative\"\nlet err_seq_max = \"Space.Sequence.create: max_length must be >= min_length\"\nlet err_text_max = \"Space.Text.create: max_length must be positive\"\nlet err_text_charset = \"Space.Text.create: charset must not be empty\"\nlet strf = Printf.sprintf\nlet errorf fmt = Format.kasprintf (fun msg -> Error msg) fmt\n\n(* Spec *)\n\ntype spec =\n | Discrete of { start : int; n : int }\n | Box of { low : float array; high : float array }\n | Multi_binary of { n : int }\n | Multi_discrete of { nvec : int array }\n | Tuple of spec list\n | Dict of (string * spec) list\n | Sequence of { min_length : int; max_length : int option; base : spec }\n | Text of { charset : string; max_length : int }\n\nlet rec equal_spec a b =\n match (a, b) with\n | Discrete a, Discrete b -> a.start = b.start && a.n = b.n\n | Box a, Box b -> a.low = b.low && a.high = b.high\n | Multi_binary a, Multi_binary b -> a.n = b.n\n | Multi_discrete a, Multi_discrete b -> a.nvec = b.nvec\n | Tuple a, Tuple b ->\n List.length a = List.length b && List.for_all2 equal_spec a b\n | Dict a, Dict b ->\n List.length a = List.length b\n && List.for_all2\n (fun (ka, sa) (kb, sb) -> String.equal ka kb && equal_spec sa sb)\n a b\n | Sequence a, Sequence b ->\n a.min_length = b.min_length\n && a.max_length = b.max_length\n && equal_spec a.base b.base\n | Text a, Text b ->\n String.equal a.charset b.charset && a.max_length = b.max_length\n | ( ( Discrete _ | Box _ | Multi_binary _ | Multi_discrete _ | Tuple _\n | Dict _ | Sequence _ | Text _ ),\n _ ) ->\n false\n\n(* Space type *)\n\ntype 'a t = {\n spec : spec;\n shape : int array option;\n contains : 'a -> bool;\n sample : unit -> 'a;\n pack : 'a -> Value.t;\n unpack : Value.t -> ('a, string) result;\n boundaries : Value.t list;\n box_bounds : (float array * float array) option;\n discrete_info : (int * int) option;\n}\n\ntype packed = Pack : 'a t -> packed\n\nlet spec s = s.spec\nlet shape s = s.shape\nlet contains s v = s.contains v\nlet sample s = s.sample ()\nlet pack s v = s.pack v\nlet unpack s v = s.unpack v\nlet boundary_values s = s.boundaries\n\n(* Discrete *)\n\nmodule Discrete = struct\n type element = (int32, Nx.int32_elt) Nx.t\n\n let to_int tensor =\n let reshaped = Nx.reshape [| 1 |] tensor in\n let arr : Int32.t array = Nx.to_array reshaped in\n Int32.to_int arr.(0)\n\n let of_int v = Nx.scalar Nx.int32 (Int32.of_int v)\n\n let create ?(start = 0) n =\n if n <= 0 then invalid_arg err_discrete_n;\n let hi = start + n in\n let contains tensor =\n let reshaped = Nx.reshape [| 1 |] tensor in\n let arr : Int32.t array = Nx.to_array reshaped in\n Array.length arr = 1\n &&\n let v = Int32.to_int arr.(0) in\n v >= start && v < hi\n in\n let sample () =\n let tensor = Nx.randint Nx.int32 ~high:hi [| 1 |] start in\n let arr : Int32.t array = Nx.to_array tensor in\n Nx.scalar Nx.int32 arr.(0)\n in\n let pack tensor =\n let arr : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] tensor) in\n Value.Int (Int32.to_int arr.(0))\n in\n let unpack = function\n | Value.Int v when v >= start && v < hi ->\n Ok (Nx.scalar Nx.int32 (Int32.of_int v))\n | Value.Int v -> errorf \"Discrete value %d outside [%d, %d)\" v start hi\n | other -> errorf \"Discrete expects Int, got %s\" (Value.to_string other)\n in\n let boundaries =\n if n = 1 then [ Value.Int start ]\n else [ Value.Int start; Value.Int (hi - 1) ]\n in\n {\n spec = Discrete { start; n };\n shape = None;\n contains;\n sample;\n pack;\n unpack;\n boundaries;\n box_bounds = None;\n discrete_info = Some (start, n);\n }\n\n let n s =\n match s.discrete_info with\n | Some (_, n) -> n\n | None -> invalid_arg err_discrete_not\n\n let start s =\n match s.discrete_info with\n | Some (start, _) -> start\n | None -> invalid_arg err_discrete_not\nend\n\n(* Box *)\n\nmodule Box = struct\n type element = (float, Nx.float32_elt) Nx.t\n\n let create ~low ~high =\n let arity = Array.length low in\n if arity = 0 then invalid_arg err_box_empty;\n if arity <> Array.length high then invalid_arg err_box_shape;\n Array.iteri\n (fun i lo ->\n if lo > high.(i) then\n invalid_arg\n (strf \"Space.Box.create: low[%d]=%g > high[%d]=%g\" i lo i high.(i)))\n low;\n let low = Array.copy low in\n let high = Array.copy high in\n let contains tensor =\n let sh = Nx.shape tensor in\n Array.length sh = 1\n && sh.(0) = arity\n &&\n let values = Nx.to_array tensor in\n let rec loop i =\n if i = arity then true\n else\n let v = values.(i) in\n v >= low.(i) && v <= high.(i) && loop (i + 1)\n in\n loop 0\n in\n let sample () =\n let uniform = Nx.rand Nx.float32 [| arity |] in\n let draws = Nx.to_array uniform in\n let values =\n Array.init arity (fun i ->\n let lo = low.(i) in\n let hi = high.(i) in\n if Float.equal lo hi then lo\n else\n let range = hi -. lo in\n if Float.is_finite range then lo +. (draws.(i) *. range)\n else\n let v = -1e6 +. (draws.(i) *. 2e6) in\n Float.max lo (Float.min hi v))\n in\n Nx.create Nx.float32 [| arity |] values\n in\n let pack tensor = Value.Float_array (Array.copy (Nx.to_array tensor)) in\n let unpack = function\n | Value.Float_array arr when Array.length arr = arity ->\n let tensor = Nx.create Nx.float32 [| arity |] arr in\n if contains tensor then Ok tensor\n else\n errorf \"Box value outside bounds: %s\"\n (Value.to_string (Value.Float_array arr))\n | Value.Float_array arr ->\n errorf \"Box expects vector of size %d, got size %d\" arity\n (Array.length arr)\n | other ->\n errorf \"Box expects Float_array, got %s\" (Value.to_string other)\n in\n let identical =\n let same = ref true in\n let i = ref 0 in\n while !same && !i < arity do\n if not (Float.equal low.(!i) high.(!i)) then same := false;\n incr i\n done;\n !same\n in\n let boundaries =\n let lo_v = Value.Float_array (Array.copy low) in\n let hi_v = Value.Float_array (Array.copy high) in\n if identical then [ lo_v ] else [ lo_v; hi_v ]\n in\n let box_bounds = Some (Array.copy low, Array.copy high) in\n {\n spec = Box { low = Array.copy low; high = Array.copy high };\n shape = Some [| arity |];\n contains;\n sample;\n pack;\n unpack;\n boundaries;\n box_bounds;\n discrete_info = None;\n }\n\n let bounds s =\n match s.box_bounds with\n | Some (low, high) -> (Array.copy low, Array.copy high)\n | None -> invalid_arg err_box_not\nend\n\n(* Multi_binary *)\n\nmodule Multi_binary = struct\n type element = (int32, Nx.int32_elt) Nx.t\n\n let create n =\n if n <= 0 then invalid_arg err_mb_n;\n let contains tensor =\n let sh = Nx.shape tensor in\n Array.length sh = 1\n && sh.(0) = n\n &&\n let arr : Int32.t array = Nx.to_array tensor in\n Array.for_all (fun v -> v = Int32.zero || v = Int32.one) arr\n in\n let sample () = Nx.randint Nx.int32 ~high:2 [| n |] 0 in\n let pack tensor =\n let arr : Int32.t array = Nx.to_array tensor in\n Value.Bool_array\n (Array.init n (fun i -> not (Int32.equal arr.(i) Int32.zero)))\n in\n let unpack = function\n | Value.Bool_array arr when Array.length arr = n ->\n let data =\n Array.map (fun b -> if b then Int32.one else Int32.zero) arr\n in\n Ok (Nx.create Nx.int32 [| n |] data)\n | Value.Bool_array arr ->\n errorf \"Multi_binary expects vector of size %d, got size %d\" n\n (Array.length arr)\n | other ->\n errorf \"Multi_binary expects Bool_array, got %s\"\n (Value.to_string other)\n in\n let boundaries =\n [\n Value.Bool_array (Array.make n false);\n Value.Bool_array (Array.make n true);\n ]\n in\n {\n spec = Multi_binary { n };\n shape = Some [| n |];\n contains;\n sample;\n pack;\n unpack;\n boundaries;\n box_bounds = None;\n discrete_info = None;\n }\nend\n\n(* Multi_discrete *)\n\nmodule Multi_discrete = struct\n type element = (int32, Nx.int32_elt) Nx.t\n\n let create nvec =\n let arity = Array.length nvec in\n if arity = 0 then invalid_arg err_md_empty;\n let nvec = Array.copy nvec in\n Array.iteri\n (fun i bound ->\n if bound <= 0 then\n invalid_arg\n (strf \"Space.Multi_discrete.create: nvec[%d] must be > 0\" i))\n nvec;\n let contains tensor =\n let sh = Nx.shape tensor in\n Array.length sh = 1\n && sh.(0) = arity\n &&\n let arr : Int32.t array = Nx.to_array tensor in\n let rec loop i =\n if i = arity then true\n else\n let v = Int32.to_int arr.(i) in\n v >= 0 && v < nvec.(i) && loop (i + 1)\n in\n loop 0\n in\n let sample () =\n let data =\n Array.init arity (fun i ->\n let tensor = Nx.randint Nx.int32 ~high:nvec.(i) [| 1 |] 0 in\n let arr = Nx.to_array tensor in\n arr.(0))\n in\n Nx.create Nx.int32 [| arity |] data\n in\n let pack tensor =\n let arr : Int32.t array = Nx.to_array tensor in\n Value.Int_array (Array.map Int32.to_int arr)\n in\n let unpack = function\n | Value.Int_array arr when Array.length arr = arity ->\n let data = Array.map Int32.of_int arr in\n let tensor = Nx.create Nx.int32 [| arity |] data in\n if contains tensor then Ok tensor\n else\n errorf \"Multi_discrete value outside bounds: %s\"\n (Value.to_string (Value.Int_array arr))\n | Value.Int_array arr ->\n errorf \"Multi_discrete expects vector of size %d, got size %d\" arity\n (Array.length arr)\n | other ->\n errorf \"Multi_discrete expects Int_array, got %s\"\n (Value.to_string other)\n in\n let boundaries =\n [\n Value.Int_array (Array.make arity 0);\n Value.Int_array (Array.init arity (fun i -> nvec.(i) - 1));\n ]\n in\n {\n spec = Multi_discrete { nvec = Array.copy nvec };\n shape = Some [| arity |];\n contains;\n sample;\n pack;\n unpack;\n boundaries;\n box_bounds = None;\n discrete_info = None;\n }\nend\n\n(* Tuple *)\n\nmodule Tuple = struct\n type element = Value.t list\n\n let create spaces =\n let spaces = Array.of_list spaces in\n let len = Array.length spaces in\n let contains values =\n let rec loop i = function\n | [] -> i = len\n | v :: rest -> (\n if i >= len then false\n else\n let (Pack s) = spaces.(i) in\n match s.unpack v with\n | Ok _ -> loop (i + 1) rest\n | Error _ -> false)\n in\n loop 0 values\n in\n let sample () =\n let values =\n Array.to_list\n (Array.init len (fun i ->\n let (Pack s) = spaces.(i) in\n let v = s.sample () in\n s.pack v))\n in\n values\n in\n let pack values = Value.List values in\n let unpack = function\n | Value.List values ->\n if List.length values <> len then\n errorf \"Tuple expects %d elements, got %d\" len (List.length values)\n else\n let rec loop i = function\n | [] -> Ok values\n | v :: rest -> (\n let (Pack s) = spaces.(i) in\n match s.unpack v with\n | Ok _ -> loop (i + 1) rest\n | Error msg -> errorf \"Tuple element %d: %s\" i msg)\n in\n loop 0 values\n | other -> errorf \"Tuple expects List, got %s\" (Value.to_string other)\n in\n let sub_specs = Array.to_list (Array.map (fun (Pack s) -> s.spec) spaces) in\n {\n spec = Tuple sub_specs;\n shape = None;\n contains;\n sample;\n pack;\n unpack;\n boundaries = [];\n box_bounds = None;\n discrete_info = None;\n }\nend\n\n(* Dict *)\n\nmodule Dict = struct\n type element = (string * Value.t) list\n\n module String_map = Map.Make (String)\n\n let create entries =\n let map =\n List.fold_left\n (fun acc (key, space) ->\n if String_map.mem key acc then\n invalid_arg (strf \"Space.Dict.create: duplicate key '%s'\" key);\n String_map.add key space acc)\n String_map.empty entries\n in\n let contains values =\n let rec loop remaining m =\n match remaining with\n | [] -> String_map.is_empty m\n | (key, value) :: rest -> (\n match String_map.find_opt key m with\n | None -> false\n | Some (Pack s) -> (\n match s.unpack value with\n | Ok _ -> loop rest (String_map.remove key m)\n | Error _ -> false))\n in\n loop values map\n in\n let sample () =\n if String_map.is_empty map then []\n else\n let acc =\n String_map.fold\n (fun key (Pack s) acc ->\n let v = s.sample () in\n (key, s.pack v) :: acc)\n map []\n in\n List.rev acc\n in\n let pack values = Value.Dict values in\n let unpack = function\n | Value.Dict values ->\n if contains values then Ok values\n else errorf \"Dict contains unexpected keys or values\"\n | other -> errorf \"Dict expects Dict, got %s\" (Value.to_string other)\n in\n let sub_specs =\n List.rev\n (String_map.fold (fun key (Pack s) acc -> (key, s.spec) :: acc) map [])\n in\n {\n spec = Dict sub_specs;\n shape = None;\n contains;\n sample;\n pack;\n unpack;\n boundaries = [];\n box_bounds = None;\n discrete_info = None;\n }\nend\n\n(* Sequence *)\n\nmodule Sequence = struct\n type 'a element = 'a list\n\n let create ?(min_length = 0) ?max_length base =\n if min_length < 0 then invalid_arg err_seq_min;\n let max_length =\n match max_length with\n | None -> None\n | Some m when m < min_length -> invalid_arg err_seq_max\n | Some _ as m -> m\n in\n let contains values =\n let len = List.length values in\n len >= min_length\n && (match max_length with None -> true | Some m -> len <= m)\n && List.for_all (fun v -> base.contains v) values\n in\n let sample () =\n let length =\n match max_length with\n | None -> min_length\n | Some max_len ->\n if max_len = min_length then min_length\n else\n let tensor =\n Nx.randint Nx.int32 ~high:(max_len + 1) [| 1 |] min_length\n in\n let arr = Nx.to_array tensor in\n Int32.to_int arr.(0)\n in\n if length = 0 then []\n else\n let rec build i acc =\n if i = length then List.rev acc\n else\n let v = base.sample () in\n build (i + 1) (v :: acc)\n in\n build 0 []\n in\n let pack values = Value.List (List.map (fun v -> base.pack v) values) in\n let unpack = function\n | Value.List values ->\n let len = List.length values in\n let exceeds =\n match max_length with None -> false | Some m -> len > m\n in\n if len < min_length || exceeds then\n match max_length with\n | None ->\n errorf \"Sequence length %d shorter than minimum %d\" len\n min_length\n | Some m ->\n errorf \"Sequence length %d outside [%d, %d]\" len min_length m\n else\n let rec loop acc = function\n | [] -> Ok (List.rev acc)\n | v :: rest -> (\n match base.unpack v with\n | Ok x -> loop (x :: acc) rest\n | Error _ as err -> err)\n in\n loop [] values\n | other -> errorf \"Sequence expects List, got %s\" (Value.to_string other)\n in\n {\n spec = Sequence { min_length; max_length; base = base.spec };\n shape = None;\n contains;\n sample;\n pack;\n unpack;\n boundaries = [];\n box_bounds = None;\n discrete_info = None;\n }\nend\n\n(* Text *)\n\nmodule Text = struct\n type element = string\n\n let default_charset =\n \"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 \"\n\n let create ?(charset = default_charset) ?(max_length = 64) () =\n if max_length <= 0 then invalid_arg err_text_max;\n let charset_len = String.length charset in\n if charset_len = 0 then invalid_arg err_text_charset;\n let contains value =\n let len = String.length value in\n len <= max_length\n &&\n let rec loop i =\n if i = len then true\n else String.contains charset value.[i] && loop (i + 1)\n in\n loop 0\n in\n let sample () =\n let length =\n if max_length = 1 then 1\n else\n let tensor = Nx.randint Nx.int32 ~high:(max_length + 1) [| 1 |] 1 in\n let arr = Nx.to_array tensor in\n Int32.to_int arr.(0)\n in\n if length = 0 then \"\"\n else\n let idxs = Nx.randint Nx.int32 ~high:charset_len [| length |] 0 in\n let arr = Nx.to_array idxs in\n Bytes.init length (fun i -> charset.[Int32.to_int arr.(i)])\n |> Bytes.to_string\n in\n let pack value = Value.String value in\n let unpack = function\n | Value.String s when contains s -> Ok s\n | Value.String s -> errorf \"Text value '%s' violates constraints\" s\n | other -> errorf \"Text expects String, got %s\" (Value.to_string other)\n in\n let example = if charset_len = 0 then \"\" else String.make 1 charset.[0] in\n let boundaries = [ Value.String \"\"; Value.String example ] in\n {\n spec = Text { charset; max_length };\n shape = None;\n contains;\n sample;\n pack;\n unpack;\n boundaries;\n box_bounds = None;\n discrete_info = None;\n }\nend\n"
},
{
"path": "packages/fehu/lib/space.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Observation and action spaces.\n\n Spaces define valid observations and actions for reinforcement learning\n environments. They specify shapes, constraints, and provide methods to\n validate, sample, and serialize values.\n\n Each space type corresponds to a common RL scenario: discrete choices,\n continuous vectors, binary indicators, composite structures, and\n variable-length sequences. *)\n\n(** {1:spec Structural description} *)\n\n(** Structural description of a space. Two spaces are compatible when their\n specs are equal. *)\ntype spec =\n | Discrete of { start : int; n : int }\n (** Integer choices in \\[[start]; [start + n - 1]\\]. *)\n | Box of { low : float array; high : float array }\n (** Continuous vector bounded per dimension. *)\n | Multi_binary of { n : int } (** Binary vector of length [n]. *)\n | Multi_discrete of { nvec : int array }\n (** Multiple discrete axes with per-axis cardinalities. *)\n | Tuple of spec list (** Fixed-length heterogeneous sequence. *)\n | Dict of (string * spec) list (** Named fields with different types. *)\n | Sequence of { min_length : int; max_length : int option; base : spec }\n (** Variable-length homogeneous sequence. *)\n | Text of { charset : string; max_length : int }\n (** Character strings from a fixed alphabet. *)\n\nval equal_spec : spec -> spec -> bool\n(** [equal_spec a b] is [true] iff [a] and [b] describe structurally identical\n spaces. *)\n\n(** {1:spaces Spaces} *)\n\ntype 'a t\n(** The type for spaces over values of type ['a]. A space is self-contained: all\n bounds, constraints, and serialization logic are stored in the value itself.\n*)\n\ntype packed =\n | Pack : 'a t -> packed\n (** Type-erased space for heterogeneous collections. *)\n\n(** {1:ops Operations} *)\n\nval spec : 'a t -> spec\n(** [spec s] is the structural description of [s]. *)\n\nval shape : 'a t -> int array option\n(** [shape s] is the dimensionality of [s], if defined. [None] for scalar or\n variable-length spaces. *)\n\nval contains : 'a t -> 'a -> bool\n(** [contains s v] is [true] iff [v] is valid in [s]. *)\n\nval sample : 'a t -> 'a\n(** [sample s] is a uniformly sampled value from [s].\n\n Random keys are drawn from the implicit RNG scope. *)\n\nval pack : 'a t -> 'a -> Value.t\n(** [pack s v] is [v] converted to the universal {!Value.t} representation. *)\n\nval unpack : 'a t -> Value.t -> ('a, string) result\n(** [unpack s v] is [Ok x] if [v] can be converted to a valid element of [s], or\n [Error msg] otherwise. *)\n\nval boundary_values : 'a t -> Value.t list\n(** [boundary_values s] is a list of representative edge-case values for [s].\n Includes lower/upper bounds or canonical sentinels when known. The empty\n list when no boundary values apply. *)\n\n(** {1:space_types Space types} *)\n\nmodule Discrete : sig\n type element = (int32, Nx.int32_elt) Nx.t\n (** Discrete action represented as a scalar int32 tensor. *)\n\n val create : ?start:int -> int -> element t\n (** [create ?start n] is a discrete space with [n] choices in the range\n \\[[start]; [start + n - 1]\\]. [start] defaults to [0].\n\n Raises [Invalid_argument] if [n <= 0]. *)\n\n val n : element t -> int\n (** [n s] is the number of choices in [s].\n\n Raises [Invalid_argument] if [s] is not a discrete space. *)\n\n val start : element t -> int\n (** [start s] is the starting value of [s].\n\n Raises [Invalid_argument] if [s] is not a discrete space. *)\n\n val to_int : element -> int\n (** [to_int e] is the integer value of the discrete element [e]. *)\n\n val of_int : int -> element\n (** [of_int v] is a discrete element with value [v]. *)\nend\n\nmodule Box : sig\n type element = (float, Nx.float32_elt) Nx.t\n (** Continuous vector represented as a float32 tensor. *)\n\n val create : low:float array -> high:float array -> element t\n (** [create ~low ~high] is a continuous space where element [i] satisfies\n [low.(i) <= x.(i) <= high.(i)]. Both arrays must have the same positive\n length.\n\n When the range of a dimension is not finite (e.g. bounds set to\n [Float.max_float]), sampling falls back to a uniform draw in \\[[-1e6];\n [1e6]\\] clamped to bounds.\n\n Raises [Invalid_argument] if [low] is empty, if [low] and [high] differ in\n length, or if any [low.(i) > high.(i)]. *)\n\n val bounds : element t -> float array * float array\n (** [bounds s] is [(low, high)] copies of the bound vectors.\n\n Raises [Invalid_argument] if [s] is not a box space. *)\nend\n\nmodule Multi_binary : sig\n type element = (int32, Nx.int32_elt) Nx.t\n (** Binary vector for multi-label scenarios. *)\n\n val create : int -> element t\n (** [create n] is a binary vector space of length [n]. Valid values are int32\n tensors with [n] elements, each 0 or 1.\n\n Raises [Invalid_argument] if [n <= 0]. *)\nend\n\nmodule Multi_discrete : sig\n type element = (int32, Nx.int32_elt) Nx.t\n (** Multiple discrete choices with independent cardinalities. *)\n\n val create : int array -> element t\n (** [create nvec] is a multi-discrete space where element [i] is in \\[[0];\n [nvec.(i) - 1]\\].\n\n Raises [Invalid_argument] if [nvec] is empty or any [nvec.(i) <= 0]. *)\nend\n\nmodule Tuple : sig\n type element = Value.t list\n (** Fixed-length heterogeneous sequence in {!Value.t} form. *)\n\n val create : packed list -> element t\n (** [create spaces] is a tuple space. Valid values are lists where element [i]\n belongs to [spaces.(i)]. {!unpack} validates each element against its\n subspace. *)\nend\n\nmodule Dict : sig\n type element = (string * Value.t) list\n (** Named fields with different space types. *)\n\n val create : (string * packed) list -> element t\n (** [create fields] is a dictionary space with named fields. Valid values are\n association lists matching the keys and subspaces of [fields].\n\n Raises [Invalid_argument] if [fields] contains duplicate keys. *)\nend\n\nmodule Sequence : sig\n type 'a element = 'a list\n (** Variable-length homogeneous sequence. *)\n\n val create : ?min_length:int -> ?max_length:int -> 'a t -> 'a element t\n (** [create ?min_length ?max_length s] is a sequence space over [s].\n [min_length] defaults to [0]. When [max_length] is provided, sampling\n draws a uniform length in \\[[min_length]; [max_length]\\]; otherwise the\n sampler returns sequences of length [min_length].\n\n Raises [Invalid_argument] if [min_length < 0] or\n [max_length < min_length]. *)\nend\n\nmodule Text : sig\n type element = string\n (** String space for textual observations or actions. *)\n\n val create : ?charset:string -> ?max_length:int -> unit -> element t\n (** [create ?charset ?max_length ()] is a text space. [charset] defaults to\n alphanumeric plus space. [max_length] defaults to [64]. Valid strings\n contain only characters from [charset] and have length at most\n [max_length].\n\n Raises [Invalid_argument] if [max_length <= 0] or [charset] is empty. *)\nend\n"
},
{
"path": "packages/fehu/lib/value.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype t =\n | Null\n | Bool of bool\n | Int of int\n | Float of float\n | String of string\n | Int_array of int array\n | Float_array of float array\n | Bool_array of bool array\n | List of t list\n | Dict of (string * t) list\n\n(* Equality *)\n\nlet rec equal a b =\n match (a, b) with\n | Null, Null -> true\n | Bool a, Bool b -> Bool.equal a b\n | Int a, Int b -> Int.equal a b\n | Float a, Float b -> Float.equal a b\n | String a, String b -> String.equal a b\n | Int_array a, Int_array b -> a = b\n | Float_array a, Float_array b -> a = b\n | Bool_array a, Bool_array b -> a = b\n | List a, List b -> equal_list a b\n | Dict a, Dict b -> equal_dict a b\n | ( ( Null | Bool _ | Int _ | Float _ | String _ | Int_array _ | Float_array _\n | Bool_array _ | List _ | Dict _ ),\n _ ) ->\n false\n\nand equal_list a b =\n match (a, b) with\n | [], [] -> true\n | x :: xs, y :: ys -> equal x y && equal_list xs ys\n | _ -> false\n\nand equal_dict a b =\n match (a, b) with\n | [], [] -> true\n | (ka, va) :: rest_a, (kb, vb) :: rest_b ->\n String.equal ka kb && equal va vb && equal_dict rest_a rest_b\n | _ -> false\n\n(* Formatting *)\n\nlet pp_array pp_elt ppf a =\n Format.fprintf ppf \"[|\";\n for i = 0 to Array.length a - 1 do\n if i > 0 then Format.fprintf ppf \"; \";\n pp_elt ppf a.(i)\n done;\n Format.fprintf ppf \"|]\"\n\nlet rec pp ppf = function\n | Null -> Format.fprintf ppf \"null\"\n | Bool b -> Format.fprintf ppf \"%b\" b\n | Int i -> Format.fprintf ppf \"%d\" i\n | Float f -> Format.fprintf ppf \"%g\" f\n | String s -> Format.fprintf ppf \"%S\" s\n | Int_array a -> pp_array (fun ppf v -> Format.fprintf ppf \"%d\" v) ppf a\n | Float_array a -> pp_array (fun ppf v -> Format.fprintf ppf \"%g\" v) ppf a\n | Bool_array a -> pp_array (fun ppf v -> Format.fprintf ppf \"%b\" v) ppf a\n | List items ->\n Format.fprintf ppf \"[\";\n List.iteri\n (fun i v ->\n if i > 0 then Format.fprintf ppf \"; \";\n pp ppf v)\n items;\n Format.fprintf ppf \"]\"\n | Dict fields ->\n Format.fprintf ppf \"{\";\n List.iteri\n (fun i (k, v) ->\n if i > 0 then Format.fprintf ppf \"; \";\n Format.fprintf ppf \"%s: %a\" k pp v)\n fields;\n Format.fprintf ppf \"}\"\n\nlet to_string v = Format.asprintf \"%a\" pp v\n"
},
{
"path": "packages/fehu/lib/value.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Universal value type.\n\n Values represent heterogeneous data flowing through spaces and info\n dictionaries. Each variant wraps one kind of scalar, array, or composite\n datum. *)\n\n(** {1:types Types} *)\n\n(** The type for universal values. *)\ntype t =\n | Null (** No value. *)\n | Bool of bool (** A boolean. *)\n | Int of int (** An integer. *)\n | Float of float (** A float. *)\n | String of string (** A string. *)\n | Int_array of int array (** An integer array. *)\n | Float_array of float array (** A float array. *)\n | Bool_array of bool array (** A boolean array. *)\n | List of t list (** A heterogeneous list. *)\n | Dict of (string * t) list (** A string-keyed association list. *)\n\n(** {1:predicates Predicates} *)\n\nval equal : t -> t -> bool\n(** [equal a b] is [true] iff [a] and [b] are structurally equal. *)\n\n(** {1:fmt Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp] formats a value for debugging. *)\n\nval to_string : t -> string\n(** [to_string v] is [v] formatted as a string via {!pp}. *)\n"
},
{
"path": "packages/fehu/lib/vec_env.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet strf = Printf.sprintf\n\n(* Error messages *)\n\nlet err_empty = \"Vec_env.create: env list must not be empty\"\nlet err_action_len n m = strf \"Vec_env.step: expected %d actions, got %d\" n m\n\nlet err_space kind =\n strf \"Vec_env.create: all environments must have the same %s space\" kind\n\n(* Types *)\n\ntype 'obs step = {\n observations : 'obs array;\n rewards : float array;\n terminated : bool array;\n truncated : bool array;\n infos : Info.t array;\n}\n\ntype ('obs, 'act, 'render) t = {\n envs : ('obs, 'act, 'render) Env.t array;\n observation_space : 'obs Space.t;\n action_space : 'act Space.t;\n}\n\n(* Space compatibility *)\n\nlet ensure_compatible envs =\n let first = envs.(0) in\n let obs_spec = Space.spec (Env.observation_space first) in\n let act_spec = Space.spec (Env.action_space first) in\n for i = 1 to Array.length envs - 1 do\n let env = envs.(i) in\n if not (Space.equal_spec obs_spec (Space.spec (Env.observation_space env)))\n then invalid_arg (err_space \"observation\");\n if not (Space.equal_spec act_spec (Space.spec (Env.action_space env))) then\n invalid_arg (err_space \"action\")\n done\n\n(* Constructor *)\n\nlet create envs =\n match envs with\n | [] -> invalid_arg err_empty\n | first :: _ ->\n let envs = Array.of_list envs in\n ensure_compatible envs;\n {\n envs;\n observation_space = Env.observation_space first;\n action_space = Env.action_space first;\n }\n\n(* Accessors *)\n\nlet num_envs t = Array.length t.envs\nlet observation_space t = t.observation_space\nlet action_space t = t.action_space\n\n(* Reset *)\n\nlet reset t () =\n let n = Array.length t.envs in\n let results = Array.init n (fun i -> Env.reset t.envs.(i) ()) in\n let observations = Array.map fst results in\n let infos = Array.map snd results in\n (observations, infos)\n\n(* Step *)\n\nlet step t actions =\n let n = Array.length t.envs in\n if Array.length actions <> n then\n invalid_arg (err_action_len n (Array.length actions));\n let results = Array.init n (fun i -> Env.step t.envs.(i) actions.(i)) in\n let observations = Array.make n results.(0).observation in\n let rewards = Array.make n 0. in\n let terminated = Array.make n false in\n let truncated = Array.make n false in\n let infos = Array.make n Info.empty in\n for i = 0 to n - 1 do\n let result = results.(i) in\n rewards.(i) <- result.reward;\n terminated.(i) <- result.terminated;\n truncated.(i) <- result.truncated;\n if result.terminated || result.truncated then begin\n let final_obs = Space.pack t.observation_space result.observation in\n let info = Info.set \"final_observation\" final_obs result.info in\n let info = Info.set \"final_info\" (Info.to_value result.info) info in\n let obs, reset_info = Env.reset t.envs.(i) () in\n observations.(i) <- obs;\n infos.(i) <- Info.merge info reset_info\n end\n else begin\n observations.(i) <- result.observation;\n infos.(i) <- result.info\n end\n done;\n { observations; rewards; terminated; truncated; infos }\n\n(* Close *)\n\nlet close t = Array.iter Env.close t.envs\n"
},
{
"path": "packages/fehu/lib/vec_env.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Vectorized environments.\n\n Runs multiple environment instances and batches their outputs. All\n environments must have compatible observation and action spaces. Terminated\n or truncated episodes are automatically reset. *)\n\n(** {1:types Types} *)\n\ntype ('obs, 'act, 'render) t\n(** The type for vectorized environments. *)\n\ntype 'obs step = {\n observations : 'obs array; (** One observation per environment. *)\n rewards : float array; (** One reward per environment. *)\n terminated : bool array; (** Per-environment termination flags. *)\n truncated : bool array; (** Per-environment truncation flags. *)\n infos : Info.t array; (** Per-environment info dictionaries. *)\n}\n(** The type for batched step results. All arrays have length {!num_envs}. *)\n\n(** {1:constructors Constructors} *)\n\nval create : ('obs, 'act, 'render) Env.t list -> ('obs, 'act, 'render) t\n(** [create envs] creates a vectorized environment.\n\n All environments must have structurally identical spaces (checked via\n {!Space.spec} and {!Space.equal_spec}). Raises [Invalid_argument] if [envs]\n is empty or spaces differ. *)\n\n(** {1:accessors Accessors} *)\n\nval num_envs : ('obs, 'act, 'render) t -> int\n(** [num_envs t] is the number of environments. *)\n\nval observation_space : ('obs, 'act, 'render) t -> 'obs Space.t\n(** [observation_space t] is the shared observation space. *)\n\nval action_space : ('obs, 'act, 'render) t -> 'act Space.t\n(** [action_space t] is the shared action space. *)\n\n(** {1:lifecycle Lifecycle} *)\n\nval reset : ('obs, 'act, 'render) t -> unit -> 'obs array * Info.t array\n(** [reset t ()] resets all environments. *)\n\nval step : ('obs, 'act, 'render) t -> 'act array -> 'obs step\n(** [step t actions] steps all environments with the given actions.\n\n [actions] must have length [num_envs t]. Terminated or truncated\n environments are automatically reset. The terminal observation is stored in\n the step's info under the key [\"final_observation\"] as a packed {!Value.t}.\n The terminal info is stored under [\"final_info\"].\n\n Raises [Invalid_argument] if [Array.length actions <> num_envs t]. *)\n\nval close : ('obs, 'act, 'render) t -> unit\n(** [close t] closes all environments. *)\n"
},
{
"path": "packages/fehu/test/dune",
"content": "(tests\n (names\n test_value\n test_info\n test_space\n test_env\n test_env_wrappers\n test_collect\n test_buffer\n test_gae\n test_eval\n test_vec_env\n test_render\n test_envs)\n (package fehu)\n (libraries fehu fehu.envs nx windtrap))\n"
},
{
"path": "packages/fehu/test/test_buffer.ml",
"content": "open Fehu\nopen Windtrap\n\nlet make_transition obs act rew next_obs term trunc =\n Buffer.\n {\n observation = obs;\n action = act;\n reward = rew;\n next_observation = next_obs;\n terminated = term;\n truncated = trunc;\n }\n\n(* Creation *)\n\nlet test_create_empty () =\n let buf = Buffer.create ~capacity:10 in\n equal ~msg:\"size = 0\" int 0 (Buffer.size buf);\n is_false ~msg:\"not full\" (Buffer.is_full buf)\n\nlet test_capacity () =\n let buf = Buffer.create ~capacity:10 in\n equal ~msg:\"capacity = 10\" int 10 (Buffer.capacity buf)\n\nlet test_create_zero_capacity () =\n raises_invalid_arg \"Buffer.create: capacity must be positive\" (fun () ->\n Buffer.create ~capacity:0)\n\nlet test_create_negative_capacity () =\n raises_invalid_arg \"Buffer.create: capacity must be positive\" (fun () ->\n Buffer.create ~capacity:(-1))\n\n(* Add/Size *)\n\nlet test_add_increments_size () =\n let buf = Buffer.create ~capacity:10 in\n Buffer.add buf (make_transition 1 0 1.0 2 false false);\n equal ~msg:\"size = 1\" int 1 (Buffer.size buf);\n Buffer.add buf (make_transition 2 1 2.0 3 false false);\n equal ~msg:\"size = 2\" int 2 (Buffer.size buf)\n\nlet test_size_capped_at_capacity () =\n let buf = Buffer.create ~capacity:3 in\n for i = 1 to 5 do\n Buffer.add buf (make_transition i 0 1.0 (i + 1) false false)\n done;\n equal ~msg:\"size capped at 3\" int 3 (Buffer.size buf)\n\nlet test_is_full () =\n let buf = Buffer.create ~capacity:2 in\n Buffer.add buf (make_transition 1 0 1.0 2 false false);\n is_false ~msg:\"not yet full\" (Buffer.is_full buf);\n Buffer.add buf (make_transition 2 1 2.0 3 false false);\n is_true ~msg:\"full\" (Buffer.is_full buf)\n\n(* Sample *)\n\nlet test_sample_batch_size () =\n let buf = Buffer.create ~capacity:10 in\n for i = 1 to 5 do\n Buffer.add buf (make_transition i 0 1.0 (i + 1) false false)\n done;\n let batch = Buffer.sample buf ~batch_size:3 in\n equal ~msg:\"batch length\" int 3 (Array.length batch)\n\nlet test_sample_empty_raises () =\n let buf = Buffer.create ~capacity:10 in\n raises_invalid_arg \"Buffer.sample: buffer is empty\" (fun () ->\n Buffer.sample buf ~batch_size:1)\n\nlet test_sample_zero_batch_raises () =\n let buf = Buffer.create ~capacity:10 in\n Buffer.add buf (make_transition 1 0 1.0 2 false false);\n raises_invalid_arg \"Buffer.sample: batch_size must be positive\" (fun () ->\n Buffer.sample buf ~batch_size:0)\n\nlet test_sample_arrays_lengths () =\n let buf = Buffer.create ~capacity:10 in\n for i = 1 to 5 do\n Buffer.add buf (make_transition i 0 1.0 (i + 1) false false)\n done;\n let obs, acts, rews, next_obs, terms, truncs =\n Buffer.sample_arrays buf ~batch_size:3\n in\n equal ~msg:\"obs length\" int 3 (Array.length obs);\n equal ~msg:\"acts length\" int 3 (Array.length acts);\n equal ~msg:\"rews length\" int 3 (Array.length rews);\n equal ~msg:\"next_obs length\" int 3 (Array.length next_obs);\n equal ~msg:\"terms length\" int 3 (Array.length terms);\n equal ~msg:\"truncs length\" int 3 (Array.length truncs)\n\nlet test_sample_arrays_empty_raises () =\n let buf = Buffer.create ~capacity:10 in\n raises_invalid_arg \"Buffer.sample: buffer is empty\" (fun () ->\n Buffer.sample_arrays buf ~batch_size:1)\n\n(* Clear *)\n\nlet test_clear_resets () =\n let buf = Buffer.create ~capacity:10 in\n Buffer.add buf (make_transition 1 0 1.0 2 false false);\n Buffer.add buf (make_transition 2 1 2.0 3 false false);\n Buffer.clear buf;\n equal ~msg:\"size = 0 after clear\" int 0 (Buffer.size buf)\n\nlet test_add_after_clear () =\n let buf = Buffer.create ~capacity:10 in\n Buffer.add buf (make_transition 1 0 1.0 2 false false);\n Buffer.clear buf;\n Buffer.add buf (make_transition 3 1 3.0 4 false false);\n equal ~msg:\"size = 1 after re-add\" int 1 (Buffer.size buf)\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu.Buffer\"\n [\n group \"creation\"\n [\n test \"empty\" test_create_empty;\n test \"capacity\" test_capacity;\n test \"zero capacity raises\" test_create_zero_capacity;\n test \"negative capacity raises\" test_create_negative_capacity;\n ];\n group \"add/size\"\n [\n test \"add increments size\" test_add_increments_size;\n test \"size capped at capacity\" test_size_capped_at_capacity;\n test \"is_full\" test_is_full;\n ];\n group \"sample\"\n [\n test \"batch size\" test_sample_batch_size;\n test \"empty raises\" test_sample_empty_raises;\n test \"zero batch raises\" test_sample_zero_batch_raises;\n test \"sample_arrays lengths\" test_sample_arrays_lengths;\n test \"sample_arrays empty raises\" test_sample_arrays_empty_raises;\n ];\n group \"clear\"\n [\n test \"resets size\" test_clear_resets;\n test \"add after clear\" test_add_after_clear;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_collect.ml",
"content": "open Fehu\nopen Windtrap\n\nlet make_test_env ?(max_steps = 100) () =\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act_space = Space.Discrete.create 2 in\n let state = ref 5.0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n state := 5.0;\n steps := 0;\n (Nx.create Nx.float32 [| 1 |] [| !state |], Info.empty)\n in\n let step _env action =\n let a : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] action) in\n state := !state +. if Int32.to_int a.(0) = 0 then -1.0 else 1.0;\n incr steps;\n let terminated = !state <= 0.0 || !state >= 10.0 in\n let truncated = (not terminated) && !steps >= max_steps in\n Env.step_result\n ~observation:(Nx.create Nx.float32 [| 1 |] [| !state |])\n ~reward:1.0 ~terminated ~truncated ()\n in\n Env.create ~id:\"Test-v0\" ~observation_space:obs_space ~action_space:act_space\n ~reset ~step ()\n\n(* Rollout *)\n\nlet test_rollout_length () =\n let env = make_test_env () in\n let policy _obs = (Nx.create Nx.int32 [| 1 |] [| 1l |], None, None) in\n let traj = Collect.rollout env ~policy ~n_steps:5 in\n equal ~msg:\"length = 5\" int 5 (Collect.length traj)\n\nlet test_rollout_arrays_length () =\n let env = make_test_env () in\n let policy _obs = (Nx.create Nx.int32 [| 1 |] [| 1l |], None, None) in\n let traj = Collect.rollout env ~policy ~n_steps:5 in\n equal ~msg:\"observations\" int 5 (Array.length traj.observations);\n equal ~msg:\"actions\" int 5 (Array.length traj.actions);\n equal ~msg:\"rewards\" int 5 (Array.length traj.rewards);\n equal ~msg:\"next_observations\" int 5 (Array.length traj.next_observations);\n equal ~msg:\"terminated\" int 5 (Array.length traj.terminated);\n equal ~msg:\"truncated\" int 5 (Array.length traj.truncated);\n equal ~msg:\"infos\" int 5 (Array.length traj.infos)\n\nlet test_rollout_next_obs_populated () =\n let env = make_test_env () in\n let policy _obs = (Nx.create Nx.int32 [| 1 |] [| 1l |], None, None) in\n let traj = Collect.rollout env ~policy ~n_steps:3 in\n for i = 0 to 2 do\n let arr : float array =\n Nx.to_array (Nx.reshape [| 1 |] traj.next_observations.(i))\n in\n is_true ~msg:\"next_obs is finite\" (Float.is_finite arr.(0))\n done\n\nlet test_rollout_no_log_probs () =\n let env = make_test_env () in\n let policy _obs = (Nx.create Nx.int32 [| 1 |] [| 1l |], None, None) in\n let traj = Collect.rollout env ~policy ~n_steps:3 in\n is_none ~msg:\"log_probs\" traj.log_probs;\n is_none ~msg:\"values\" traj.values\n\nlet test_rollout_with_log_probs () =\n let env = make_test_env () in\n let policy _obs =\n (Nx.create Nx.int32 [| 1 |] [| 1l |], Some (-0.5), Some 1.0)\n in\n let traj = Collect.rollout env ~policy ~n_steps:4 in\n is_some ~msg:\"log_probs present\" traj.log_probs;\n is_some ~msg:\"values present\" traj.values;\n equal ~msg:\"log_probs length\" int 4 (Array.length (Option.get traj.log_probs));\n equal ~msg:\"values length\" int 4 (Array.length (Option.get traj.values))\n\n(* Episodes *)\n\nlet test_episodes_count () =\n let env = make_test_env ~max_steps:10 () in\n let policy _obs = (Nx.create Nx.int32 [| 1 |] [| 1l |], None, None) in\n let eps = Collect.episodes env ~policy ~n_episodes:2 ~max_steps:10 () in\n equal ~msg:\"2 episodes\" int 2 (List.length eps)\n\nlet test_episodes_positive_length () =\n let env = make_test_env ~max_steps:10 () in\n let policy _obs = (Nx.create Nx.int32 [| 1 |] [| 1l |], None, None) in\n let eps = Collect.episodes env ~policy ~n_episodes:2 ~max_steps:10 () in\n List.iter\n (fun ep ->\n is_true ~msg:\"episode has positive length\" (Collect.length ep > 0))\n eps\n\n(* Concat *)\n\nlet test_concat_two () =\n let env = make_test_env () in\n let policy _obs = (Nx.create Nx.int32 [| 1 |] [| 1l |], None, None) in\n let t1 = Collect.rollout env ~policy ~n_steps:3 in\n let t2 = Collect.rollout env ~policy ~n_steps:4 in\n let t = Collect.concat [ t1; t2 ] in\n equal ~msg:\"total length\" int 7 (Collect.length t)\n\nlet test_concat_empty_raises () =\n raises_invalid_arg \"Collect.concat: empty list\" (fun () -> Collect.concat [])\n\nlet test_concat_singleton () =\n let env = make_test_env () in\n let policy _obs = (Nx.create Nx.int32 [| 1 |] [| 1l |], None, None) in\n let t1 = Collect.rollout env ~policy ~n_steps:5 in\n let t = Collect.concat [ t1 ] in\n equal ~msg:\"same length\" int 5 (Collect.length t)\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu.Collect\"\n [\n group \"rollout\"\n [\n test \"length\" test_rollout_length;\n test \"arrays length\" test_rollout_arrays_length;\n test \"next_observations populated\" test_rollout_next_obs_populated;\n test \"no log_probs/values\" test_rollout_no_log_probs;\n test \"with log_probs/values\" test_rollout_with_log_probs;\n ];\n group \"episodes\"\n [\n test \"count\" test_episodes_count;\n test \"positive length\" test_episodes_positive_length;\n ];\n group \"concat\"\n [\n test \"two trajectories\" test_concat_two;\n test \"empty raises\" test_concat_empty_raises;\n test \"singleton\" test_concat_singleton;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_env.ml",
"content": "open Fehu\nopen Windtrap\n\nlet make_test_env ?(max_steps = 100) () =\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act_space = Space.Discrete.create 2 in\n let state = ref 5.0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n state := 5.0;\n steps := 0;\n (Nx.create Nx.float32 [| 1 |] [| !state |], Info.empty)\n in\n let step _env action =\n let a : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] action) in\n state := !state +. if Int32.to_int a.(0) = 0 then -1.0 else 1.0;\n incr steps;\n let terminated = !state <= 0.0 || !state >= 10.0 in\n let truncated = (not terminated) && !steps >= max_steps in\n Env.step_result\n ~observation:(Nx.create Nx.float32 [| 1 |] [| !state |])\n ~reward:1.0 ~terminated ~truncated ()\n in\n Env.create ~id:\"Test-v0\" ~observation_space:obs_space ~action_space:act_space\n ~reset ~step ()\n\nlet action_left = Nx.create Nx.int32 [| 1 |] [| 0l |]\nlet action_right = Nx.create Nx.int32 [| 1 |] [| 1l |]\n\nlet read_obs obs =\n let arr : float array = Nx.to_array (Nx.reshape [| 1 |] obs) in\n arr.(0)\n\n(* Creation *)\n\nlet test_id () =\n let env = make_test_env () in\n equal ~msg:\"id is Some Test-v0\" (option string) (Some \"Test-v0\") (Env.id env)\n\nlet test_observation_space () =\n let env = make_test_env () in\n let low, high = Space.Box.bounds (Env.observation_space env) in\n equal ~msg:\"obs low\" (array (float 0.0)) [| 0.0 |] low;\n equal ~msg:\"obs high\" (array (float 0.0)) [| 10.0 |] high\n\nlet test_action_space () =\n let env = make_test_env () in\n equal ~msg:\"act n\" int 2 (Space.Discrete.n (Env.action_space env))\n\nlet test_render_mode_default () =\n let env = make_test_env () in\n is_none ~msg:\"render_mode default is None\" (Env.render_mode env)\n\nlet test_render_mode_invalid () =\n raises_invalid_arg ~msg:\"render_mode not in render_modes\"\n \"Env.create: render mode 'human' not in render_modes []\" (fun () ->\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n let act_space = Space.Discrete.create 2 in\n Env.create ~observation_space:obs_space ~action_space:act_space\n ~render_mode:`Human ~render_modes:[]\n ~reset:(fun _env ?options:_ () -> assert false)\n ~step:(fun _env _ -> assert false)\n ())\n\n(* Lifecycle *)\n\nlet test_reset_obs () =\n let env = make_test_env () in\n let obs, _info = Env.reset env () in\n equal ~msg:\"reset obs shape\" (array int) [| 1 |] (Nx.shape obs);\n equal ~msg:\"reset obs value\" (float 0.0) 5.0 (read_obs obs)\n\nlet test_step_after_reset () =\n let env = make_test_env () in\n let _obs, _info = Env.reset env () in\n let step = Env.step env action_right in\n equal ~msg:\"reward\" (float 0.0) 1.0 step.reward;\n is_false ~msg:\"not terminated\" step.terminated;\n is_false ~msg:\"not truncated\" step.truncated\n\nlet test_step_before_reset () =\n let env = make_test_env () in\n raises_invalid_arg ~msg:\"step before reset\"\n \"Env: operation 'step' requires calling reset first\" (fun () ->\n Env.step env action_left)\n\nlet test_step_after_terminal () =\n let env = make_test_env () in\n let _obs, _info = Env.reset env () in\n (* Move left 5 times: 5 -> 4 -> 3 -> 2 -> 1 -> 0, terminated *)\n for _ = 1 to 4 do\n ignore (Env.step env action_left)\n done;\n let step = Env.step env action_left in\n is_true ~msg:\"terminated\" step.terminated;\n raises_invalid_arg ~msg:\"step after terminal\"\n \"Env: operation 'step' requires calling reset first\" (fun () ->\n Env.step env action_left)\n\nlet test_reset_after_terminal () =\n let env = make_test_env () in\n let _obs, _info = Env.reset env () in\n for _ = 1 to 5 do\n ignore (Env.step env action_left)\n done;\n let obs, _info = Env.reset env () in\n equal ~msg:\"reset clears terminal\" (float 0.0) 5.0 (read_obs obs)\n\nlet test_close () =\n let env = make_test_env () in\n Env.close env;\n is_true ~msg:\"closed\" (Env.closed env)\n\nlet test_step_on_closed () =\n let env = make_test_env () in\n let _obs, _info = Env.reset env () in\n Env.close env;\n raises_invalid_arg ~msg:\"step on closed\"\n \"Env: operation 'step' on a closed environment\" (fun () ->\n Env.step env action_left)\n\nlet test_reset_on_closed () =\n let env = make_test_env () in\n Env.close env;\n raises_invalid_arg ~msg:\"reset on closed\"\n \"Env: operation 'reset' on a closed environment\" (fun () ->\n Env.reset env ())\n\nlet test_render_on_closed () =\n let env = make_test_env () in\n Env.close env;\n raises_invalid_arg ~msg:\"render on closed\"\n \"Env: operation 'render' on a closed environment\" (fun () -> Env.render env)\n\nlet test_close_idempotent () =\n let env = make_test_env () in\n Env.close env;\n Env.close env;\n is_true ~msg:\"still closed\" (Env.closed env)\n\n(* step_result *)\n\nlet test_step_result_defaults () =\n let obs = Nx.create Nx.float32 [| 1 |] [| 0.0 |] in\n let s = Env.step_result ~observation:obs () in\n equal ~msg:\"default reward\" (float 0.0) 0.0 s.reward;\n is_false ~msg:\"default terminated\" s.terminated;\n is_false ~msg:\"default truncated\" s.truncated;\n is_true ~msg:\"default info empty\" (Info.is_empty s.info)\n\nlet test_step_result_custom () =\n let obs = Nx.create Nx.float32 [| 1 |] [| 0.0 |] in\n let info = Info.set \"k\" (Info.int 1) Info.empty in\n let s =\n Env.step_result ~observation:obs ~reward:5.0 ~terminated:true\n ~truncated:false ~info ()\n in\n equal ~msg:\"custom reward\" (float 0.0) 5.0 s.reward;\n is_true ~msg:\"custom terminated\" s.terminated;\n is_false ~msg:\"custom truncated\" s.truncated;\n is_some ~msg:\"custom info has key\" (Info.find \"k\" s.info)\n\n(* time_limit lifecycle enforcement *)\n\nlet test_time_limit_needs_reset () =\n let env = make_test_env () in\n let wrapped = Env.time_limit ~max_episode_steps:3 env in\n let _obs, _info = Env.reset wrapped () in\n for _ = 1 to 2 do\n ignore (Env.step wrapped action_right)\n done;\n let s3 = Env.step wrapped action_right in\n is_true ~msg:\"step 3 truncated\" s3.truncated;\n raises_invalid_arg ~msg:\"step after time_limit truncation requires reset\"\n \"Env: operation 'step' requires calling reset first\" (fun () ->\n Env.step wrapped action_right)\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu.Env\"\n [\n group \"creation\"\n [\n test \"id\" test_id;\n test \"observation_space\" test_observation_space;\n test \"action_space\" test_action_space;\n test \"render_mode default\" test_render_mode_default;\n test \"render_mode invalid\" test_render_mode_invalid;\n ];\n group \"lifecycle\"\n [\n test \"reset returns valid obs\" test_reset_obs;\n test \"step after reset\" test_step_after_reset;\n test \"step before reset\" test_step_before_reset;\n test \"step after terminal\" test_step_after_terminal;\n test \"reset after terminal\" test_reset_after_terminal;\n test \"close\" test_close;\n test \"step on closed\" test_step_on_closed;\n test \"reset on closed\" test_reset_on_closed;\n test \"render on closed\" test_render_on_closed;\n test \"close idempotent\" test_close_idempotent;\n test \"time_limit needs reset after truncation\"\n test_time_limit_needs_reset;\n ];\n group \"step_result\"\n [\n test \"defaults\" test_step_result_defaults;\n test \"custom values\" test_step_result_custom;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_env_wrappers.ml",
"content": "open Fehu\nopen Windtrap\n\nlet make_test_env ?(max_steps = 100) () =\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act_space = Space.Discrete.create 2 in\n let state = ref 5.0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n state := 5.0;\n steps := 0;\n (Nx.create Nx.float32 [| 1 |] [| !state |], Info.empty)\n in\n let step _env action =\n let a : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] action) in\n state := !state +. if Int32.to_int a.(0) = 0 then -1.0 else 1.0;\n incr steps;\n let terminated = !state <= 0.0 || !state >= 10.0 in\n let truncated = (not terminated) && !steps >= max_steps in\n Env.step_result\n ~observation:(Nx.create Nx.float32 [| 1 |] [| !state |])\n ~reward:1.0 ~terminated ~truncated ()\n in\n Env.create ~id:\"Test-v0\" ~observation_space:obs_space ~action_space:act_space\n ~reset ~step ()\n\nlet action_left = Nx.create Nx.int32 [| 1 |] [| 0l |]\nlet action_right = Nx.create Nx.int32 [| 1 |] [| 1l |]\n\nlet read_obs obs =\n let arr : float array = Nx.to_array (Nx.reshape [| 1 |] obs) in\n arr.(0)\n\nlet value = testable ~pp:Value.pp ~equal:Value.equal ()\n\n(* State sharing *)\n\nlet test_close_wrapper_closes_inner () =\n let env = make_test_env () in\n let wrapped =\n Env.map_observation\n ~observation_space:(Env.observation_space env)\n ~f:(fun obs info -> (obs, info))\n env\n in\n Env.close wrapped;\n is_true ~msg:\"inner closed\" (Env.closed env)\n\nlet test_close_inner_closes_wrapper () =\n let env = make_test_env () in\n let wrapped =\n Env.map_observation\n ~observation_space:(Env.observation_space env)\n ~f:(fun obs info -> (obs, info))\n env\n in\n Env.close env;\n is_true ~msg:\"wrapper closed\" (Env.closed wrapped)\n\nlet test_reset_wrapper_clears_inner () =\n let env = make_test_env () in\n let wrapped =\n Env.map_observation\n ~observation_space:(Env.observation_space env)\n ~f:(fun obs info -> (obs, info))\n env\n in\n let _obs, _info = Env.reset wrapped () in\n let step = Env.step env action_left in\n equal ~msg:\"inner step works\" (float 0.0) 1.0 step.reward\n\n(* map_observation *)\n\nlet test_map_observation_reset () =\n let env = make_test_env () in\n let double_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 20.0 |] in\n let wrapped =\n Env.map_observation ~observation_space:double_space\n ~f:(fun obs info ->\n let v = read_obs obs in\n (Nx.create Nx.float32 [| 1 |] [| v *. 2.0 |], info))\n env\n in\n let obs, _info = Env.reset wrapped () in\n equal ~msg:\"doubled reset obs\" (float 0.0) 10.0 (read_obs obs)\n\nlet test_map_observation_step () =\n let env = make_test_env () in\n let double_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 20.0 |] in\n let wrapped =\n Env.map_observation ~observation_space:double_space\n ~f:(fun obs info ->\n let v = read_obs obs in\n (Nx.create Nx.float32 [| 1 |] [| v *. 2.0 |], info))\n env\n in\n let _obs, _info = Env.reset wrapped () in\n let step = Env.step wrapped action_right in\n (* Inner: 5 + 1 = 6, doubled: 12 *)\n equal ~msg:\"doubled step obs\" (float 0.0) 12.0 (read_obs step.observation)\n\nlet test_map_observation_id () =\n let env = make_test_env () in\n let wrapped =\n Env.map_observation\n ~observation_space:(Env.observation_space env)\n ~f:(fun obs info -> (obs, info))\n env\n in\n equal ~msg:\"id suffix\" (option string) (Some \"Test-v0/ObservationWrapper\")\n (Env.id wrapped)\n\n(* map_action *)\n\nlet test_map_action_flip () =\n let env = make_test_env () in\n let wrapped =\n Env.map_action ~action_space:(Env.action_space env)\n ~f:(fun action ->\n let a : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] action) in\n let flipped = if Int32.to_int a.(0) = 0 then 1l else 0l in\n Nx.create Nx.int32 [| 1 |] [| flipped |])\n env\n in\n let _obs, _info = Env.reset wrapped () in\n (* Send left (0) to wrapper; inner sees right (1): 5 -> 6 *)\n let step = Env.step wrapped action_left in\n equal ~msg:\"flipped: left becomes right\" (float 0.0) 6.0\n (read_obs step.observation)\n\nlet test_map_action_id () =\n let env = make_test_env () in\n let wrapped =\n Env.map_action ~action_space:(Env.action_space env) ~f:Fun.id env\n in\n equal ~msg:\"id suffix\" (option string) (Some \"Test-v0/ActionWrapper\")\n (Env.id wrapped)\n\n(* map_reward *)\n\nlet test_map_reward () =\n let env = make_test_env () in\n let wrapped =\n Env.map_reward ~f:(fun ~reward ~info -> (reward *. 2.0, info)) env\n in\n let _obs, _info = Env.reset wrapped () in\n let step = Env.step wrapped action_right in\n equal ~msg:\"doubled reward\" (float 0.0) 2.0 step.reward\n\nlet test_map_reward_id () =\n let env = make_test_env () in\n let wrapped = Env.map_reward ~f:(fun ~reward ~info -> (reward, info)) env in\n equal ~msg:\"id suffix\" (option string) (Some \"Test-v0/RewardWrapper\")\n (Env.id wrapped)\n\n(* clip_action *)\n\nlet make_box_action_env () =\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n let last_action = ref 0.0 in\n let reset _env ?options:_ () =\n last_action := 0.0;\n (Nx.create Nx.float32 [| 1 |] [| 5.0 |], Info.empty)\n in\n let step _env action =\n let a : float array = Nx.to_array (Nx.reshape [| 1 |] action) in\n last_action := a.(0);\n Env.step_result\n ~observation:(Nx.create Nx.float32 [| 1 |] [| a.(0) |])\n ~reward:1.0 ()\n in\n let env =\n Env.create ~id:\"BoxAct-v0\" ~observation_space:obs_space\n ~action_space:act_space ~reset ~step ()\n in\n (env, last_action)\n\nlet test_clip_action () =\n let env, last_action = make_box_action_env () in\n let wrapped = Env.clip_action env in\n let _obs, _info = Env.reset wrapped () in\n let _step = Env.step wrapped (Nx.create Nx.float32 [| 1 |] [| 2.0 |]) in\n equal ~msg:\"clamped to upper\" (float 0.0) 1.0 !last_action;\n let _step = Env.step wrapped (Nx.create Nx.float32 [| 1 |] [| -0.5 |]) in\n equal ~msg:\"clamped to lower\" (float 0.0) 0.0 !last_action\n\n(* clip_observation *)\n\nlet make_box_obs_env () =\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act_space = Space.Discrete.create 2 in\n let obs_val = ref 5.0 in\n let reset _env ?options:_ () =\n obs_val := 5.0;\n (Nx.create Nx.float32 [| 1 |] [| !obs_val |], Info.empty)\n in\n let step _env action =\n let a : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] action) in\n obs_val := !obs_val +. if Int32.to_int a.(0) = 0 then -3.0 else 3.0;\n Env.step_result\n ~observation:(Nx.create Nx.float32 [| 1 |] [| !obs_val |])\n ~reward:1.0 ()\n in\n Env.create ~id:\"BoxObs-v0\" ~observation_space:obs_space\n ~action_space:act_space ~reset ~step ()\n\nlet test_clip_observation () =\n let env = make_box_obs_env () in\n let wrapped = Env.clip_observation ~low:[| 2.0 |] ~high:[| 8.0 |] env in\n let _obs, _info = Env.reset wrapped () in\n (* Step right: inner obs = 8.0, clipped to 8.0 *)\n let s1 = Env.step wrapped action_right in\n let arr1 : float array = Nx.to_array (Nx.reshape [| 1 |] s1.observation) in\n equal ~msg:\"clipped to upper\" (float 0.0) 8.0 arr1.(0);\n let _obs, _info = Env.reset wrapped () in\n (* Step left: inner obs = 2.0, within bounds *)\n let s2 = Env.step wrapped action_left in\n let arr2 : float array = Nx.to_array (Nx.reshape [| 1 |] s2.observation) in\n equal ~msg:\"within bounds\" (float 0.0) 2.0 arr2.(0)\n\nlet test_clip_observation_space () =\n let env = make_box_obs_env () in\n let wrapped = Env.clip_observation ~low:[| 2.0 |] ~high:[| 8.0 |] env in\n let low, high = Space.Box.bounds (Env.observation_space wrapped) in\n equal ~msg:\"clipped low\" (array (float 0.0)) [| 2.0 |] low;\n equal ~msg:\"clipped high\" (array (float 0.0)) [| 8.0 |] high\n\n(* time_limit *)\n\nlet test_time_limit_truncation () =\n let env = make_test_env () in\n let wrapped = Env.time_limit ~max_episode_steps:3 env in\n let _obs, _info = Env.reset wrapped () in\n let s1 = Env.step wrapped action_right in\n is_false ~msg:\"step 1 not truncated\" s1.truncated;\n let s2 = Env.step wrapped action_right in\n is_false ~msg:\"step 2 not truncated\" s2.truncated;\n let s3 = Env.step wrapped action_right in\n is_true ~msg:\"step 3 truncated\" s3.truncated\n\nlet test_time_limit_info () =\n let env = make_test_env () in\n let wrapped = Env.time_limit ~max_episode_steps:2 env in\n let _obs, _info = Env.reset wrapped () in\n let _s1 = Env.step wrapped action_right in\n let s2 = Env.step wrapped action_right in\n is_some ~msg:\"time_limit.truncated present\"\n (Info.find \"time_limit.truncated\" s2.info);\n is_some ~msg:\"time_limit.elapsed_steps present\"\n (Info.find \"time_limit.elapsed_steps\" s2.info)\n\nlet test_time_limit_info_values () =\n let env = make_test_env () in\n let wrapped = Env.time_limit ~max_episode_steps:2 env in\n let _obs, _info = Env.reset wrapped () in\n let _s1 = Env.step wrapped action_right in\n let s2 = Env.step wrapped action_right in\n let tl_truncated = Info.find_exn \"time_limit.truncated\" s2.info in\n equal ~msg:\"truncated is Bool true\" value (Value.Bool true) tl_truncated;\n let tl_steps = Info.find_exn \"time_limit.elapsed_steps\" s2.info in\n equal ~msg:\"elapsed_steps is Int 2\" value (Value.Int 2) tl_steps\n\nlet test_time_limit_counter_resets () =\n let env = make_test_env () in\n let wrapped = Env.time_limit ~max_episode_steps:3 env in\n let _obs, _info = Env.reset wrapped () in\n for _ = 1 to 3 do\n ignore (Env.step wrapped action_right)\n done;\n let _obs, _info = Env.reset wrapped () in\n let s1 = Env.step wrapped action_right in\n is_false ~msg:\"counter reset after new episode\" s1.truncated\n\nlet test_time_limit_nonpositive () =\n let env = make_test_env () in\n raises_invalid_arg ~msg:\"max_episode_steps=0\"\n \"Env.time_limit: max_episode_steps must be positive\" (fun () ->\n Env.time_limit ~max_episode_steps:0 env);\n raises_invalid_arg ~msg:\"max_episode_steps=-1\"\n \"Env.time_limit: max_episode_steps must be positive\" (fun () ->\n Env.time_limit ~max_episode_steps:(-1) env)\n\nlet test_time_limit_needs_reset () =\n let env = make_test_env () in\n let wrapped = Env.time_limit ~max_episode_steps:2 env in\n let _obs, _info = Env.reset wrapped () in\n let _s1 = Env.step wrapped action_right in\n let s2 = Env.step wrapped action_right in\n is_true ~msg:\"truncated at limit\" s2.truncated;\n raises_invalid_arg ~msg:\"step after time_limit truncation\"\n \"Env: operation 'step' requires calling reset first\" (fun () ->\n Env.step wrapped action_right)\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu.Env (wrappers)\"\n [\n group \"state sharing\"\n [\n test \"close wrapper closes inner\" test_close_wrapper_closes_inner;\n test \"close inner closes wrapper\" test_close_inner_closes_wrapper;\n test \"reset wrapper clears inner\" test_reset_wrapper_clears_inner;\n ];\n group \"map_observation\"\n [\n test \"doubles reset obs\" test_map_observation_reset;\n test \"doubles step obs\" test_map_observation_step;\n test \"id suffix\" test_map_observation_id;\n ];\n group \"map_action\"\n [\n test \"flip reverses direction\" test_map_action_flip;\n test \"id suffix\" test_map_action_id;\n ];\n group \"map_reward\"\n [\n test \"doubles reward\" test_map_reward;\n test \"id suffix\" test_map_reward_id;\n ];\n group \"clip_action\" [ test \"clamps out-of-bounds\" test_clip_action ];\n group \"clip_observation\"\n [\n test \"clamps to explicit bounds\" test_clip_observation;\n test \"observation space reflects bounds\" test_clip_observation_space;\n ];\n group \"time_limit\"\n [\n test \"truncation at limit\" test_time_limit_truncation;\n test \"info keys present\" test_time_limit_info;\n test \"info values correct\" test_time_limit_info_values;\n test \"counter resets on new episode\" test_time_limit_counter_resets;\n test \"nonpositive raises\" test_time_limit_nonpositive;\n test \"needs reset after truncation\" test_time_limit_needs_reset;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_envs.ml",
"content": "open Fehu\nopen Fehu_envs\nopen Windtrap\n\nlet read_float obs =\n let arr : float array = Nx.to_array (Nx.reshape [| 1 |] obs) in\n arr.(0)\n\nlet read_int32_array obs n =\n let arr : Int32.t array = Nx.to_array (Nx.reshape [| n |] obs) in\n Array.map Int32.to_int arr\n\nlet discrete action = Nx.create Nx.int32 [| 1 |] [| Int32.of_int action |]\n\n(* Random_walk *)\n\nlet test_rw_creation () =\n let env = Random_walk.make () in\n match Env.id env with\n | Some id ->\n is_true ~msg:\"id starts with RandomWalk\"\n (String.length id >= 10 && String.sub id 0 10 = \"RandomWalk\")\n | None -> fail \"expected an id\"\n\nlet test_rw_reset_obs () =\n let env = Random_walk.make () in\n let obs, _info = Env.reset env () in\n equal ~msg:\"reset obs is 0.0\" (float 1e-6) 0.0 (read_float obs)\n\nlet test_rw_step_left () =\n let env = Random_walk.make () in\n let _obs, _info = Env.reset env () in\n let s = Env.step env (discrete 0) in\n equal ~msg:\"step left to -1.0\" (float 1e-6) (-1.0) (read_float s.observation)\n\nlet test_rw_step_right () =\n let env = Random_walk.make () in\n let _obs, _info = Env.reset env () in\n let s = Env.step env (discrete 1) in\n equal ~msg:\"step right to 1.0\" (float 1e-6) 1.0 (read_float s.observation)\n\nlet test_rw_termination () =\n let env = Random_walk.make () in\n let _obs, _info = Env.reset env () in\n let terminated = ref false in\n for _ = 1 to 20 do\n if not !terminated then begin\n let s = Env.step env (discrete 1) in\n if s.terminated then terminated := true\n else if s.truncated then begin\n let _obs, _info = Env.reset env () in\n ()\n end\n end\n done;\n is_true ~msg:\"terminated at boundary\" !terminated\n\nlet test_rw_ansi_render () =\n let env = Random_walk.make ~render_mode:`Ansi () in\n let _obs, _info = Env.reset env () in\n match Env.render env with\n | Some s -> is_true ~msg:\"non-empty render\" (String.length s > 0)\n | None -> fail \"expected Some render\"\n\n(* Cartpole *)\n\nlet test_cp_creation () =\n let env = Cartpole.make () in\n match Env.id env with\n | Some id ->\n is_true ~msg:\"id starts with CartPole\"\n (String.length id >= 8 && String.sub id 0 8 = \"CartPole\")\n | None -> fail \"expected an id\"\n\nlet test_cp_reset_shape () =\n let env = Cartpole.make () in\n let obs, _info = Env.reset env () in\n let shape = Nx.shape obs in\n equal ~msg:\"obs shape [4]\" (array int) [| 4 |] shape\n\nlet test_cp_step_reward () =\n let env = Cartpole.make () in\n let _obs, _info = Env.reset env () in\n let s = Env.step env (discrete 1) in\n is_false ~msg:\"not terminated on first step\" s.terminated;\n equal ~msg:\"reward 1.0\" (float 1e-6) 1.0 s.reward\n\nlet test_cp_termination () =\n let env = Cartpole.make () in\n let _obs, _info = Env.reset env () in\n let done_flag = ref false in\n for _ = 1 to 600 do\n if not !done_flag then begin\n let s = Env.step env (discrete 0) in\n if s.terminated || s.truncated then done_flag := true\n end\n done;\n is_true ~msg:\"episode ends\" !done_flag\n\n(* Grid_world *)\n\nlet test_gw_creation () =\n let env = Grid_world.make () in\n match Env.id env with\n | Some id ->\n is_true ~msg:\"id starts with GridWorld\"\n (String.length id >= 9 && String.sub id 0 9 = \"GridWorld\")\n | None -> fail \"expected an id\"\n\nlet test_gw_reset_obs () =\n let env = Grid_world.make () in\n let obs, _info = Env.reset env () in\n let pos = read_int32_array obs 2 in\n equal ~msg:\"row = 0\" int 0 pos.(0);\n equal ~msg:\"col = 0\" int 0 pos.(1)\n\nlet test_gw_move_down () =\n let env = Grid_world.make () in\n let _obs, _info = Env.reset env () in\n let s = Env.step env (discrete 1) in\n let pos = read_int32_array s.observation 2 in\n equal ~msg:\"row = 1 after down\" int 1 pos.(0)\n\nlet test_gw_move_right () =\n let env = Grid_world.make () in\n let _obs, _info = Env.reset env () in\n let s = Env.step env (discrete 3) in\n let pos = read_int32_array s.observation 2 in\n equal ~msg:\"col = 1 after right\" int 1 pos.(1)\n\nlet test_gw_obstacle () =\n let env = Grid_world.make () in\n let _obs, _info = Env.reset env () in\n (* Navigate to (1, 2): down, right, right *)\n let _s = Env.step env (discrete 1) in\n let _s = Env.step env (discrete 3) in\n let s = Env.step env (discrete 3) in\n let pos = read_int32_array s.observation 2 in\n equal ~msg:\"at (1,2)\" int 1 pos.(0);\n equal ~msg:\"at (1,2)\" int 2 pos.(1);\n (* Try to move down into obstacle at (2,2) *)\n let s = Env.step env (discrete 1) in\n let pos = read_int32_array s.observation 2 in\n equal ~msg:\"blocked row still 1\" int 1 pos.(0);\n equal ~msg:\"blocked col still 2\" int 2 pos.(1)\n\nlet test_gw_reach_goal () =\n let env = Grid_world.make () in\n let _obs, _info = Env.reset env () in\n (* Path to (4,4) avoiding obstacle at (2,2): down 4 times to row 4, then right\n 4 times to col 4 *)\n for _ = 1 to 4 do\n ignore (Env.step env (discrete 1))\n done;\n let s_right1 = Env.step env (discrete 3) in\n is_false ~msg:\"not done yet\" s_right1.terminated;\n let _s = Env.step env (discrete 3) in\n let _s = Env.step env (discrete 3) in\n let s = Env.step env (discrete 3) in\n is_true ~msg:\"terminated at goal\" s.terminated;\n equal ~msg:\"reward 10.0\" (float 1e-6) 10.0 s.reward\n\nlet test_gw_ansi_render () =\n let env = Grid_world.make ~render_mode:`Ansi () in\n let _obs, _info = Env.reset env () in\n match Env.render env with\n | Some (Grid_world.Text s) ->\n is_true ~msg:\"non-empty render\" (String.length s > 0)\n | Some (Grid_world.Image _) -> fail \"expected Text render\"\n | None -> fail \"expected Some render\"\n\n(* Mountain_car *)\n\nlet test_mc_creation () =\n let env = Mountain_car.make () in\n match Env.id env with\n | Some id ->\n is_true ~msg:\"id starts with MountainCar\"\n (String.length id >= 11 && String.sub id 0 11 = \"MountainCar\")\n | None -> fail \"expected an id\"\n\nlet test_mc_reset_shape () =\n let env = Mountain_car.make () in\n let obs, _info = Env.reset env () in\n let shape = Nx.shape obs in\n equal ~msg:\"obs shape [2]\" (array int) [| 2 |] shape\n\nlet test_mc_step_coast () =\n let env = Mountain_car.make () in\n let _obs, _info = Env.reset env () in\n let s = Env.step env (discrete 1) in\n let shape = Nx.shape s.observation in\n equal ~msg:\"obs shape after step\" (array int) [| 2 |] shape;\n is_false ~msg:\"not terminated\" s.terminated\n\nlet test_mc_reward () =\n let env = Mountain_car.make () in\n let _obs, _info = Env.reset env () in\n let s = Env.step env (discrete 1) in\n equal ~msg:\"reward -1.0\" (float 1e-6) (-1.0) s.reward\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu_envs\"\n [\n group \"RandomWalk\"\n [\n test \"creation\" test_rw_creation;\n test \"reset observation\" test_rw_reset_obs;\n test \"step left\" test_rw_step_left;\n test \"step right\" test_rw_step_right;\n test \"termination at boundary\" test_rw_termination;\n test \"ansi render\" test_rw_ansi_render;\n ];\n group \"CartPole\"\n [\n test \"creation\" test_cp_creation;\n test \"reset shape\" test_cp_reset_shape;\n test \"step reward\" test_cp_step_reward;\n test \"termination\" test_cp_termination;\n ];\n group \"GridWorld\"\n [\n test \"creation\" test_gw_creation;\n test \"reset observation\" test_gw_reset_obs;\n test \"move down\" test_gw_move_down;\n test \"move right\" test_gw_move_right;\n test \"obstacle blocks movement\" test_gw_obstacle;\n test \"reach goal\" test_gw_reach_goal;\n test \"ansi render\" test_gw_ansi_render;\n ];\n group \"MountainCar\"\n [\n test \"creation\" test_mc_creation;\n test \"reset shape\" test_mc_reset_shape;\n test \"step coast\" test_mc_step_coast;\n test \"reward\" test_mc_reward;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_eval.ml",
"content": "open Fehu\nopen Windtrap\n\nlet make_test_env ?(max_steps = 100) () =\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act_space = Space.Discrete.create 2 in\n let state = ref 5.0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n state := 5.0;\n steps := 0;\n (Nx.create Nx.float32 [| 1 |] [| !state |], Info.empty)\n in\n let step _env action =\n let a : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] action) in\n state := !state +. if Int32.to_int a.(0) = 0 then -1.0 else 1.0;\n incr steps;\n let terminated = !state <= 0.0 || !state >= 10.0 in\n let truncated = (not terminated) && !steps >= max_steps in\n Env.step_result\n ~observation:(Nx.create Nx.float32 [| 1 |] [| !state |])\n ~reward:1.0 ~terminated ~truncated ()\n in\n Env.create ~id:\"Test-v0\" ~observation_space:obs_space ~action_space:act_space\n ~reset ~step ()\n\n(* Run *)\n\nlet test_constant_reward_stats () =\n let env = make_test_env ~max_steps:5 () in\n let policy _obs = Nx.create Nx.int32 [| 1 |] [| 1l |] in\n let stats = Eval.run env ~policy ~n_episodes:3 ~max_steps:5 () in\n equal ~msg:\"mean_reward\" (float 1e-6) 5.0 stats.mean_reward;\n equal ~msg:\"std_reward\" (float 1e-6) 0.0 stats.std_reward;\n equal ~msg:\"mean_length\" (float 1e-6) 5.0 stats.mean_length;\n equal ~msg:\"n_episodes\" int 3 stats.n_episodes\n\nlet test_n_episodes_matches () =\n let env = make_test_env ~max_steps:5 () in\n let policy _obs = Nx.create Nx.int32 [| 1 |] [| 1l |] in\n let stats = Eval.run env ~policy ~n_episodes:7 ~max_steps:5 () in\n equal ~msg:\"n_episodes matches\" int 7 stats.n_episodes\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu.Eval\"\n [\n group \"run\"\n [\n test \"constant reward statistics\" test_constant_reward_stats;\n test \"n_episodes matches stats\" test_n_episodes_matches;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_gae.ml",
"content": "open Fehu\nopen Windtrap\n\nlet f = float 1e-6\n\n(* Compute *)\n\nlet test_compute_simple () =\n let rewards = [| 1.0; 1.0; 1.0 |] in\n let values = [| 0.5; 0.5; 0.5 |] in\n let terminated = [| false; false; false |] in\n let truncated = [| false; false; false |] in\n let next_values = [| 0.5; 0.5; 0.5 |] in\n let advantages, returns =\n Gae.compute ~rewards ~values ~terminated ~truncated ~next_values ~gamma:0.99\n ~lambda:0.95\n in\n equal ~msg:\"lengths match\" int 3 (Array.length advantages);\n for i = 0 to 2 do\n equal ~msg:\"returns = advantages + values\" f returns.(i)\n (advantages.(i) +. values.(i))\n done\n\nlet test_compute_termination () =\n let rewards = [| 1.0; 1.0; 1.0 |] in\n let values = [| 0.5; 0.5; 0.5 |] in\n let terminated = [| false; true; false |] in\n let truncated = [| false; false; false |] in\n let next_values = [| 0.5; 0.5; 0.5 |] in\n let advantages, _returns =\n Gae.compute ~rewards ~values ~terminated ~truncated ~next_values ~gamma:0.99\n ~lambda:0.95\n in\n (* At step 1 (terminated), bootstrap is 0: delta = 1.0 + 0.99*0 - 0.5 = 0.5 *)\n equal ~msg:\"terminal advantage\" f 0.5 advantages.(1)\n\nlet test_compute_truncation () =\n let rewards = [| 1.0; 1.0; 1.0 |] in\n let values = [| 0.5; 0.5; 0.5 |] in\n let terminated = [| false; false; false |] in\n let truncated = [| false; true; false |] in\n let next_values = [| 0.5; 2.0; 0.5 |] in\n let advantages_trunc, _returns =\n Gae.compute ~rewards ~values ~terminated ~truncated ~next_values ~gamma:0.99\n ~lambda:0.95\n in\n (* At step 1 (truncated), bootstrap uses next_values.(1) = 2.0 delta = 1.0 +\n 0.99*2.0 - 0.5 = 2.48 *)\n let terminated_fake = [| false; true; false |] in\n let advantages_term, _returns_term =\n Gae.compute ~rewards ~values ~terminated:terminated_fake\n ~truncated:[| false; false; false |] ~next_values ~gamma:0.99 ~lambda:0.95\n in\n (* With termination instead, bootstrap would be 0: delta = 1.0 + 0.99*0 - 0.5\n = 0.5 These must differ because truncation uses next_values. *)\n not_equal ~msg:\"truncation differs from termination\" f advantages_trunc.(1)\n advantages_term.(1)\n\nlet test_compute_length_mismatch () =\n raises_invalid_arg \"Gae: all arrays must have the same length\" (fun () ->\n Gae.compute ~rewards:[| 1.0; 1.0 |] ~values:[| 0.5 |]\n ~terminated:[| false; false |] ~truncated:[| false; false |]\n ~next_values:[| 0.5; 0.5 |] ~gamma:0.99 ~lambda:0.95)\n\nlet test_compute_empty () =\n let advantages, returns =\n Gae.compute ~rewards:[||] ~values:[||] ~terminated:[||] ~truncated:[||]\n ~next_values:[||] ~gamma:0.99 ~lambda:0.95\n in\n equal ~msg:\"empty advantages\" int 0 (Array.length advantages);\n equal ~msg:\"empty returns\" int 0 (Array.length returns)\n\n(* Returns *)\n\nlet test_returns_simple () =\n let ret =\n Gae.returns ~rewards:[| 1.0; 1.0; 1.0 |]\n ~terminated:[| false; false; false |] ~truncated:[| false; false; false |]\n ~gamma:1.0\n in\n equal ~msg:\"ret[0]\" f 3.0 ret.(0);\n equal ~msg:\"ret[1]\" f 2.0 ret.(1);\n equal ~msg:\"ret[2]\" f 1.0 ret.(2)\n\nlet test_returns_gamma_zero () =\n let ret =\n Gae.returns ~rewards:[| 1.0; 2.0; 3.0 |]\n ~terminated:[| false; false; false |] ~truncated:[| false; false; false |]\n ~gamma:0.0\n in\n equal ~msg:\"ret[0]\" f 1.0 ret.(0);\n equal ~msg:\"ret[1]\" f 2.0 ret.(1);\n equal ~msg:\"ret[2]\" f 3.0 ret.(2)\n\nlet test_returns_terminated () =\n let ret =\n Gae.returns ~rewards:[| 1.0; 1.0; 1.0 |]\n ~terminated:[| false; true; false |] ~truncated:[| false; false; false |]\n ~gamma:1.0\n in\n (* Step 2: acc = 1.0 Step 1: terminated, so acc = 1.0 + 1.0*0.0*1.0 = 1.0 Step\n 0: acc = 1.0 + 1.0*1.0*1.0 = 2.0 *)\n equal ~msg:\"ret[0]\" f 2.0 ret.(0);\n equal ~msg:\"ret[1]\" f 1.0 ret.(1);\n equal ~msg:\"ret[2]\" f 1.0 ret.(2)\n\nlet test_returns_truncated () =\n let ret =\n Gae.returns ~rewards:[| 1.0; 1.0; 1.0 |]\n ~terminated:[| false; false; false |] ~truncated:[| false; true; false |]\n ~gamma:1.0\n in\n (* Truncation at step 1 resets accumulation, same as terminated *)\n equal ~msg:\"ret[0]\" f 2.0 ret.(0);\n equal ~msg:\"ret[1]\" f 1.0 ret.(1);\n equal ~msg:\"ret[2]\" f 1.0 ret.(2)\n\nlet test_returns_length_mismatch () =\n raises_invalid_arg\n \"Gae.returns: rewards, terminated, and truncated must have the same length\"\n (fun () ->\n Gae.returns ~rewards:[| 1.0; 1.0 |] ~terminated:[| false |]\n ~truncated:[| false; false |] ~gamma:0.99)\n\n(* Compute from values *)\n\nlet test_compute_from_values_simple () =\n let rewards = [| 1.0; 1.0; 1.0 |] in\n let values = [| 0.5; 0.5; 0.5 |] in\n let terminated = [| false; false; false |] in\n let truncated = [| false; false; false |] in\n let last_value = 0.5 in\n let advantages, returns =\n Gae.compute_from_values ~rewards ~values ~terminated ~truncated ~last_value\n ~gamma:0.99 ~lambda:0.95\n in\n (* next_values should be [| 0.5; 0.5; 0.5 |] (values shifted + last_value) *)\n let advantages2, returns2 =\n Gae.compute ~rewards ~values ~terminated ~truncated\n ~next_values:[| 0.5; 0.5; 0.5 |] ~gamma:0.99 ~lambda:0.95\n in\n for i = 0 to 2 do\n equal ~msg:\"advantages match\" f advantages2.(i) advantages.(i);\n equal ~msg:\"returns match\" f returns2.(i) returns.(i)\n done\n\nlet test_compute_from_values_shifted () =\n let rewards = [| 1.0; 1.0; 1.0 |] in\n let values = [| 1.0; 2.0; 3.0 |] in\n let terminated = [| false; false; false |] in\n let truncated = [| false; false; false |] in\n let last_value = 4.0 in\n let advantages, _returns =\n Gae.compute_from_values ~rewards ~values ~terminated ~truncated ~last_value\n ~gamma:0.99 ~lambda:0.95\n in\n (* next_values = [| 2.0; 3.0; 4.0 |] *)\n let advantages2, _returns2 =\n Gae.compute ~rewards ~values ~terminated ~truncated\n ~next_values:[| 2.0; 3.0; 4.0 |] ~gamma:0.99 ~lambda:0.95\n in\n for i = 0 to 2 do\n equal ~msg:\"advantages match\" f advantages2.(i) advantages.(i)\n done\n\n(* Normalize *)\n\nlet test_normalize_mean_zero () =\n let arr = [| 1.0; 2.0; 3.0; 4.0; 5.0 |] in\n let normed = Gae.normalize arr in\n let mean = ref 0.0 in\n Array.iter (fun x -> mean := !mean +. x) normed;\n mean := !mean /. Float.of_int (Array.length normed);\n equal ~msg:\"mean near 0\" f 0.0 !mean\n\nlet test_normalize_std_one () =\n let arr = [| 1.0; 2.0; 3.0; 4.0; 5.0 |] in\n let normed = Gae.normalize arr in\n let n = Array.length normed in\n let mean = ref 0.0 in\n Array.iter (fun x -> mean := !mean +. x) normed;\n mean := !mean /. Float.of_int n;\n let var = ref 0.0 in\n Array.iter\n (fun x ->\n let d = x -. !mean in\n var := !var +. (d *. d))\n normed;\n var := !var /. Float.of_int n;\n let std = sqrt !var in\n is_true ~msg:\"std near 1\" (Float.abs (std -. 1.0) < 0.01)\n\nlet test_normalize_empty () =\n let normed = Gae.normalize [||] in\n equal ~msg:\"empty\" int 0 (Array.length normed)\n\nlet test_normalize_single () =\n let normed = Gae.normalize [| 42.0 |] in\n equal ~msg:\"single normalizes to 0\" f 0.0 normed.(0)\n\nlet () =\n run \"Fehu.Gae\"\n [\n group \"compute\"\n [\n test \"simple\" test_compute_simple;\n test \"termination\" test_compute_termination;\n test \"truncation\" test_compute_truncation;\n test \"length mismatch\" test_compute_length_mismatch;\n test \"empty\" test_compute_empty;\n ];\n group \"returns\"\n [\n test \"simple\" test_returns_simple;\n test \"gamma zero\" test_returns_gamma_zero;\n test \"terminated resets\" test_returns_terminated;\n test \"truncated resets\" test_returns_truncated;\n test \"length mismatch\" test_returns_length_mismatch;\n ];\n group \"compute_from_values\"\n [\n test \"matches compute\" test_compute_from_values_simple;\n test \"shifted values\" test_compute_from_values_shifted;\n ];\n group \"normalize\"\n [\n test \"mean near zero\" test_normalize_mean_zero;\n test \"std near one\" test_normalize_std_one;\n test \"empty\" test_normalize_empty;\n test \"single element\" test_normalize_single;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_info.ml",
"content": "open Fehu\nopen Windtrap\n\nlet value = testable ~pp:Value.pp ~equal:Value.equal ()\n\n(* Operations *)\n\nlet test_empty_is_empty () =\n equal ~msg:\"empty is_empty\" bool true (Info.is_empty Info.empty)\n\nlet test_set_then_find () =\n let info = Info.set \"k\" (Value.Int 1) Info.empty in\n equal ~msg:\"find after set\" (option value) (Some (Value.Int 1))\n (Info.find \"k\" info)\n\nlet test_find_missing () =\n let info = Info.set \"k\" (Value.Int 1) Info.empty in\n equal ~msg:\"find missing\" (option value) None (Info.find \"other\" info)\n\nlet test_find_exn_existing () =\n let info = Info.set \"k\" (Value.Int 42) Info.empty in\n equal ~msg:\"find_exn existing\" value (Value.Int 42) (Info.find_exn \"k\" info)\n\nlet test_remove () =\n let info = Info.set \"k\" (Value.Int 1) Info.empty in\n let info = Info.remove \"k\" info in\n equal ~msg:\"find after remove\" (option value) None (Info.find \"k\" info)\n\nlet test_merge_right_biased () =\n let a = Info.set \"k\" (Value.Int 1) Info.empty in\n let b = Info.set \"k\" (Value.Int 2) Info.empty in\n let merged = Info.merge a b in\n equal ~msg:\"merge right wins\" value (Value.Int 2) (Info.find_exn \"k\" merged)\n\nlet test_of_list_to_list_round_trip () =\n let kvs = [ (\"a\", Value.Int 1); (\"c\", Value.Int 3); (\"b\", Value.Int 2) ] in\n let info = Info.of_list kvs in\n let result = Info.to_list info in\n equal ~msg:\"round-trip keys sorted\" (list string) [ \"a\"; \"b\"; \"c\" ]\n (List.map fst result);\n equal ~msg:\"round-trip values\" (list value)\n [ Value.Int 1; Value.Int 2; Value.Int 3 ]\n (List.map snd result)\n\n(* Errors *)\n\nlet test_find_exn_missing () =\n raises_invalid_arg \"Info.find_exn: key \\\"missing\\\" not present\" (fun () ->\n ignore (Info.find_exn \"missing\" Info.empty))\n\n(* Convenience *)\n\nlet test_int_convenience () =\n equal ~msg:\"Info.int\" value (Value.Int 42) (Info.int 42)\n\nlet test_float_convenience () =\n equal ~msg:\"Info.float\" value (Value.Float 1.0) (Info.float 1.0)\n\nlet test_bool_convenience () =\n equal ~msg:\"Info.bool\" value (Value.Bool true) (Info.bool true)\n\nlet test_string_convenience () =\n equal ~msg:\"Info.string\" value (Value.String \"hi\") (Info.string \"hi\")\n\nlet test_null_convenience () = equal ~msg:\"Info.null\" value Value.Null Info.null\n\nlet () =\n run \"Fehu.Info\"\n [\n group \"operations\"\n [\n test \"empty is_empty\" test_empty_is_empty;\n test \"set then find\" test_set_then_find;\n test \"find missing key\" test_find_missing;\n test \"find_exn existing\" test_find_exn_existing;\n test \"remove\" test_remove;\n test \"merge right-biased\" test_merge_right_biased;\n test \"of_list/to_list round-trip\" test_of_list_to_list_round_trip;\n ];\n group \"errors\" [ test \"find_exn missing raises\" test_find_exn_missing ];\n group \"convenience\"\n [\n test \"int\" test_int_convenience;\n test \"float\" test_float_convenience;\n test \"bool\" test_bool_convenience;\n test \"string\" test_string_convenience;\n test \"null\" test_null_convenience;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_render.ml",
"content": "open Fehu\nopen Windtrap\n\nlet make_data n =\n Bigarray.Array1.create Bigarray.int8_unsigned Bigarray.c_layout n\n\n(* Image *)\n\nlet test_valid_rgb_image () =\n let data = make_data 12 in\n let img = Render.image ~width:2 ~height:2 data in\n equal ~msg:\"width\" int 2 img.width;\n equal ~msg:\"height\" int 2 img.height\n\nlet test_wrong_data_length_raises () =\n let data = make_data 10 in\n raises_invalid_arg\n \"Render.image: data length 10 does not match width * height * channels = 12\"\n (fun () -> ignore (Render.image ~width:2 ~height:2 data))\n\nlet test_rgba_channels () =\n let data = make_data 16 in\n let img =\n Render.image ~width:2 ~height:2 ~pixel_format:Render.Pixel.Rgba data\n in\n equal ~msg:\"width\" int 2 img.width;\n equal ~msg:\"height\" int 2 img.height\n\nlet test_gray_channels () =\n let data = make_data 4 in\n let img =\n Render.image ~width:2 ~height:2 ~pixel_format:Render.Pixel.Gray data\n in\n equal ~msg:\"width\" int 2 img.width;\n equal ~msg:\"height\" int 2 img.height\n\nlet test_pixel_format_default_rgb () =\n let data = make_data 3 in\n let img = Render.image ~width:1 ~height:1 data in\n equal ~msg:\"default is Rgb\" int 3 (Render.Pixel.channels img.pixel_format)\n\n(* Rollout *)\n\nlet make_renderable_env () =\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act_space = Space.Discrete.create 2 in\n let state = ref 5.0 in\n let reset _env ?options:_ () =\n state := 5.0;\n (Nx.create Nx.float32 [| 1 |] [| !state |], Info.empty)\n in\n let step _env action =\n let a : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] action) in\n state := !state +. if Int32.to_int a.(0) = 0 then -1.0 else 1.0;\n let terminated = !state <= 0.0 || !state >= 10.0 in\n Env.step_result\n ~observation:(Nx.create Nx.float32 [| 1 |] [| !state |])\n ~reward:1.0 ~terminated ()\n in\n let render () =\n let data = make_data 3 in\n Some (Render.image ~width:1 ~height:1 data)\n in\n Env.create ~id:\"Renderable-v0\" ~observation_space:obs_space\n ~action_space:act_space ~reset ~step ~render ()\n\nlet test_rollout_sink_called () =\n let env = make_renderable_env () in\n let count = ref 0 in\n let policy _obs = Nx.create Nx.int32 [| 1 |] [| 1l |] in\n let sink _frame = incr count in\n Render.rollout env ~policy ~steps:3 ~sink ();\n equal ~msg:\"sink called 3 times\" int 3 !count\n\n(* on_render *)\n\nlet action_right = Nx.create Nx.int32 [| 1 |] [| 1l |]\n\nlet test_on_render_frame_count () =\n let env = make_renderable_env () in\n let count = ref 0 in\n let wrapped = Render.on_render ~sink:(fun _ -> incr count) env in\n let _obs, _info = Env.reset wrapped () in\n let _s1 = Env.step wrapped action_right in\n let _s2 = Env.step wrapped action_right in\n let _s3 = Env.step wrapped action_right in\n (* 1 frame from reset + 3 frames from steps = 4 *)\n equal ~msg:\"frame count\" int 4 !count\n\nlet test_on_render_passthrough () =\n let env = make_renderable_env () in\n let wrapped = Render.on_render ~sink:(fun _ -> ()) env in\n let _obs, _info = Env.reset wrapped () in\n let step = Env.step wrapped action_right in\n equal ~msg:\"reward unchanged\" (float 0.0) 1.0 step.reward;\n is_false ~msg:\"not terminated\" step.terminated;\n is_false ~msg:\"not truncated\" step.truncated\n\nlet test_on_render_id () =\n let env = make_renderable_env () in\n let wrapped = Render.on_render ~sink:(fun _ -> ()) env in\n equal ~msg:\"id suffix\" (option string) (Some \"Renderable-v0/OnRender\")\n (Env.id wrapped)\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu.Render\"\n [\n group \"image\"\n [\n test \"valid RGB 2x2\" test_valid_rgb_image;\n test \"wrong data length raises\" test_wrong_data_length_raises;\n test \"RGBA 4 channels\" test_rgba_channels;\n test \"Gray 1 channel\" test_gray_channels;\n test \"default pixel_format is Rgb\" test_pixel_format_default_rgb;\n ];\n group \"rollout\"\n [ test \"sink called for each step\" test_rollout_sink_called ];\n group \"on_render\"\n [\n test \"frame count\" test_on_render_frame_count;\n test \"passthrough\" test_on_render_passthrough;\n test \"id suffix\" test_on_render_id;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_space.ml",
"content": "open Fehu\nopen Windtrap\n\nlet value = testable ~pp:Value.pp ~equal:Value.equal ()\n\n(* Helpers *)\n\nlet int32_scalar v = Nx.scalar Nx.int32 (Int32.of_int v)\n\nlet int32_vec arr =\n Nx.create Nx.int32 [| Array.length arr |] (Array.map Int32.of_int arr)\n\nlet float32_vec arr = Nx.create Nx.float32 [| Array.length arr |] arr\n\nlet read_float32_vec t =\n let n = (Nx.shape t).(0) in\n let arr : float array = Nx.to_array (Nx.reshape [| n |] t) in\n arr\n\n(* Discrete *)\n\nlet test_discrete_default () =\n let s = Space.Discrete.create 3 in\n equal ~msg:\"n is 3\" int 3 (Space.Discrete.n s);\n equal ~msg:\"start is 0\" int 0 (Space.Discrete.start s)\n\nlet test_discrete_custom_start () =\n let s = Space.Discrete.create ~start:5 3 in\n equal ~msg:\"start is 5\" int 5 (Space.Discrete.start s);\n equal ~msg:\"n is 3\" int 3 (Space.Discrete.n s)\n\nlet test_discrete_contains () =\n let s = Space.Discrete.create 3 in\n is_true ~msg:\"contains 0\" (Space.contains s (int32_scalar 0));\n is_true ~msg:\"contains 1\" (Space.contains s (int32_scalar 1));\n is_true ~msg:\"contains 2\" (Space.contains s (int32_scalar 2));\n is_false ~msg:\"not contains 3\" (Space.contains s (int32_scalar 3));\n is_false ~msg:\"not contains -1\" (Space.contains s (int32_scalar (-1)))\n\nlet test_discrete_contains_with_start () =\n let s = Space.Discrete.create ~start:5 3 in\n is_true ~msg:\"contains 5\" (Space.contains s (int32_scalar 5));\n is_true ~msg:\"contains 7\" (Space.contains s (int32_scalar 7));\n is_false ~msg:\"not contains 4\" (Space.contains s (int32_scalar 4));\n is_false ~msg:\"not contains 8\" (Space.contains s (int32_scalar 8))\n\nlet test_discrete_sample () =\n let s = Space.Discrete.create 3 in\n let v = Space.sample s in\n is_true ~msg:\"sample is valid\" (Space.contains s v)\n\nlet test_discrete_pack_unpack () =\n let s = Space.Discrete.create 3 in\n let v = int32_scalar 2 in\n let packed = Space.pack s v in\n equal ~msg:\"pack produces Int 2\" value (Value.Int 2) packed;\n let unpacked = Space.unpack s packed in\n is_ok ~msg:\"unpack succeeds\" unpacked\n\nlet test_discrete_unpack_invalid () =\n let s = Space.Discrete.create 3 in\n is_error ~msg:\"unpack out of range\" (Space.unpack s (Value.Int 5));\n is_error ~msg:\"unpack wrong type\" (Space.unpack s (Value.String \"x\"))\n\nlet test_discrete_boundary_values () =\n let s = Space.Discrete.create 3 in\n let bvs = Space.boundary_values s in\n equal ~msg:\"2 boundary values\" int 2 (List.length bvs);\n equal ~msg:\"first boundary\" value (Value.Int 0) (List.hd bvs);\n equal ~msg:\"last boundary\" value (Value.Int 2) (List.nth bvs 1)\n\nlet test_discrete_boundary_single () =\n let s = Space.Discrete.create 1 in\n let bvs = Space.boundary_values s in\n equal ~msg:\"1 boundary for n=1\" int 1 (List.length bvs)\n\nlet test_discrete_shape () =\n let s = Space.Discrete.create 3 in\n is_none ~msg:\"discrete shape is None\" (Space.shape s)\n\nlet test_discrete_error_zero () =\n raises_invalid_arg \"Space.Discrete.create: n must be strictly positive\"\n (fun () -> Space.Discrete.create 0)\n\nlet test_discrete_error_negative () =\n raises_invalid_arg \"Space.Discrete.create: n must be strictly positive\"\n (fun () -> Space.Discrete.create (-1))\n\n(* Box *)\n\nlet test_box_1d () =\n let s = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let low, high = Space.Box.bounds s in\n equal ~msg:\"low\" (array (float 0.)) [| 0.0 |] low;\n equal ~msg:\"high\" (array (float 0.)) [| 10.0 |] high\n\nlet test_box_contains () =\n let s = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n is_true ~msg:\"mid value\" (Space.contains s (float32_vec [| 5.0 |]));\n is_true ~msg:\"low bound\" (Space.contains s (float32_vec [| 0.0 |]));\n is_true ~msg:\"high bound\" (Space.contains s (float32_vec [| 10.0 |]));\n is_false ~msg:\"below low\" (Space.contains s (float32_vec [| -0.1 |]));\n is_false ~msg:\"above high\" (Space.contains s (float32_vec [| 10.1 |]))\n\nlet test_box_sample () =\n let s = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let v = Space.sample s in\n is_true ~msg:\"sample in bounds\" (Space.contains s v)\n\nlet test_box_pack_unpack () =\n let s = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let v = float32_vec [| 5.0 |] in\n let packed = Space.pack s v in\n let unpacked = Space.unpack s packed in\n is_ok ~msg:\"round-trip succeeds\" unpacked;\n match unpacked with\n | Ok t ->\n let arr = read_float32_vec t in\n equal ~msg:\"value preserved\" (float 0.01) 5.0 arr.(0)\n | Error _ -> fail \"unreachable\"\n\nlet test_box_boundary_values () =\n let s = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let bvs = Space.boundary_values s in\n equal ~msg:\"2 boundaries\" int 2 (List.length bvs)\n\nlet test_box_boundary_identical () =\n let s = Space.Box.create ~low:[| 5.0 |] ~high:[| 5.0 |] in\n let bvs = Space.boundary_values s in\n equal ~msg:\"1 boundary when identical\" int 1 (List.length bvs)\n\nlet test_box_shape_1d () =\n let s = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n is_some ~msg:\"shape is Some\" (Space.shape s);\n equal ~msg:\"shape [|1|]\" (array int) [| 1 |] (Option.get (Space.shape s))\n\nlet test_box_2d () =\n let s = Space.Box.create ~low:[| 0.0; -1.0 |] ~high:[| 1.0; 1.0 |] in\n equal ~msg:\"shape [|2|]\" (array int) [| 2 |] (Option.get (Space.shape s));\n is_true ~msg:\"2d in bounds\" (Space.contains s (float32_vec [| 0.5; 0.0 |]));\n is_false ~msg:\"2d out of bounds\"\n (Space.contains s (float32_vec [| 0.5; 2.0 |]))\n\nlet test_box_error_empty () =\n raises_invalid_arg \"Space.Box.create: low cannot be empty\" (fun () ->\n Space.Box.create ~low:[||] ~high:[||])\n\nlet test_box_error_mismatch () =\n raises_invalid_arg\n \"Space.Box.create: low and high must have identical lengths\" (fun () ->\n Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0; 2.0 |])\n\nlet test_box_error_low_gt_high () =\n raises_match ~msg:\"low > high raises\"\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> Space.Box.create ~low:[| 5.0 |] ~high:[| 1.0 |])\n\n(* Multi_binary *)\n\nlet test_mb_contains () =\n let s = Space.Multi_binary.create 3 in\n is_true ~msg:\"all zeros\" (Space.contains s (int32_vec [| 0; 0; 0 |]));\n is_true ~msg:\"all ones\" (Space.contains s (int32_vec [| 1; 1; 1 |]));\n is_true ~msg:\"mixed\" (Space.contains s (int32_vec [| 0; 1; 0 |]));\n is_false ~msg:\"value 2 invalid\" (Space.contains s (int32_vec [| 0; 2; 0 |]));\n is_false ~msg:\"wrong length\" (Space.contains s (int32_vec [| 0; 1 |]))\n\nlet test_mb_sample () =\n let s = Space.Multi_binary.create 3 in\n let v = Space.sample s in\n is_true ~msg:\"sample valid\" (Space.contains s v)\n\nlet test_mb_boundary_values () =\n let s = Space.Multi_binary.create 3 in\n let bvs = Space.boundary_values s in\n equal ~msg:\"2 boundaries\" int 2 (List.length bvs)\n\nlet test_mb_shape () =\n let s = Space.Multi_binary.create 3 in\n equal ~msg:\"shape [|3|]\" (option (array int)) (Some [| 3 |]) (Space.shape s)\n\nlet test_mb_error () =\n raises_invalid_arg \"Space.Multi_binary.create: n must be strictly positive\"\n (fun () -> Space.Multi_binary.create 0)\n\n(* Multi_discrete *)\n\nlet test_md_contains () =\n let s = Space.Multi_discrete.create [| 3; 4 |] in\n is_true ~msg:\"valid\" (Space.contains s (int32_vec [| 0; 0 |]));\n is_true ~msg:\"upper valid\" (Space.contains s (int32_vec [| 2; 3 |]));\n is_false ~msg:\"first oob\" (Space.contains s (int32_vec [| 3; 0 |]));\n is_false ~msg:\"second oob\" (Space.contains s (int32_vec [| 0; 4 |]));\n is_false ~msg:\"negative\" (Space.contains s (int32_vec [| -1; 0 |]))\n\nlet test_md_sample () =\n let s = Space.Multi_discrete.create [| 3; 4 |] in\n let v = Space.sample s in\n is_true ~msg:\"sample valid\" (Space.contains s v)\n\nlet test_md_shape () =\n let s = Space.Multi_discrete.create [| 3; 4 |] in\n equal ~msg:\"shape [|2|]\" (option (array int)) (Some [| 2 |]) (Space.shape s)\n\nlet test_md_error_empty () =\n raises_invalid_arg \"Space.Multi_discrete.create: nvec must not be empty\"\n (fun () -> Space.Multi_discrete.create [||])\n\nlet test_md_error_zero_element () =\n raises_match ~msg:\"nvec element <= 0 raises\"\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> Space.Multi_discrete.create [| 3; 0 |])\n\n(* Tuple *)\n\nlet test_tuple_contains () =\n let ds = Space.Discrete.create 3 in\n let bs = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n let s = Space.Tuple.create [ Pack ds; Pack bs ] in\n let valid = [ Value.Int 1; Value.Float_array [| 0.5 |] ] in\n is_true ~msg:\"valid tuple\" (Space.contains s valid);\n let bad_length = [ Value.Int 1 ] in\n is_false ~msg:\"wrong length\" (Space.contains s bad_length);\n let bad_value = [ Value.Int 5; Value.Float_array [| 0.5 |] ] in\n is_false ~msg:\"invalid element\" (Space.contains s bad_value)\n\nlet test_tuple_sample () =\n let ds = Space.Discrete.create 3 in\n let bs = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n let s = Space.Tuple.create [ Pack ds; Pack bs ] in\n let v = Space.sample s in\n is_true ~msg:\"sample valid\" (Space.contains s v)\n\nlet test_tuple_pack_unpack () =\n let ds = Space.Discrete.create 3 in\n let bs = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n let s = Space.Tuple.create [ Pack ds; Pack bs ] in\n let v = [ Value.Int 1; Value.Float_array [| 0.5 |] ] in\n let packed = Space.pack s v in\n let unpacked = Space.unpack s packed in\n is_ok ~msg:\"round-trip succeeds\" unpacked\n\nlet test_tuple_empty () =\n let s = Space.Tuple.create [] in\n is_true ~msg:\"empty tuple valid\" (Space.contains s []);\n is_false ~msg:\"non-empty invalid\" (Space.contains s [ Value.Int 0 ])\n\n(* Dict *)\n\nlet test_dict_contains () =\n let ds = Space.Discrete.create 3 in\n let bs = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n let s = Space.Dict.create [ (\"action\", Pack ds); (\"obs\", Pack bs) ] in\n let valid =\n [ (\"action\", Value.Int 1); (\"obs\", Value.Float_array [| 0.5 |]) ]\n in\n is_true ~msg:\"valid dict\" (Space.contains s valid);\n let missing_key = [ (\"action\", Value.Int 1) ] in\n is_false ~msg:\"missing key\" (Space.contains s missing_key);\n let extra_key =\n [\n (\"action\", Value.Int 1);\n (\"obs\", Value.Float_array [| 0.5 |]);\n (\"extra\", Value.Int 0);\n ]\n in\n is_false ~msg:\"extra key\" (Space.contains s extra_key)\n\nlet test_dict_sample () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Dict.create [ (\"a\", Pack ds) ] in\n let v = Space.sample s in\n is_true ~msg:\"sample valid\" (Space.contains s v)\n\nlet test_dict_error_duplicate () =\n let ds = Space.Discrete.create 3 in\n raises_match ~msg:\"duplicate key raises\"\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> Space.Dict.create [ (\"a\", Pack ds); (\"a\", Pack ds) ])\n\n(* Text *)\n\nlet test_text_contains () =\n let s = Space.Text.create () in\n is_true ~msg:\"alpha string\" (Space.contains s \"hello\");\n is_true ~msg:\"empty string\" (Space.contains s \"\");\n is_true ~msg:\"with digits\" (Space.contains s \"abc123\");\n is_true ~msg:\"with space\" (Space.contains s \"hello world\")\n\nlet test_text_contains_invalid () =\n let s = Space.Text.create ~charset:\"abc\" () in\n is_false ~msg:\"char outside charset\" (Space.contains s \"abcd\")\n\nlet test_text_contains_too_long () =\n let s = Space.Text.create ~max_length:3 () in\n is_false ~msg:\"exceeds max_length\" (Space.contains s \"abcd\");\n is_true ~msg:\"at max_length\" (Space.contains s \"abc\")\n\nlet test_text_sample () =\n let s = Space.Text.create () in\n let v = Space.sample s in\n is_true ~msg:\"sample valid\" (Space.contains s v);\n is_true ~msg:\"sample non-empty\" (String.length v > 0)\n\nlet test_text_boundary_values () =\n let s = Space.Text.create () in\n let bvs = Space.boundary_values s in\n equal ~msg:\"2 boundaries\" int 2 (List.length bvs)\n\nlet test_text_error_max_length () =\n raises_invalid_arg \"Space.Text.create: max_length must be positive\" (fun () ->\n Space.Text.create ~max_length:0 ())\n\nlet test_text_error_charset () =\n raises_invalid_arg \"Space.Text.create: charset must not be empty\" (fun () ->\n Space.Text.create ~charset:\"\" ())\n\n(* Sequence *)\n\nlet test_seq_contains () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Sequence.create ~min_length:1 ~max_length:3 ds in\n let v1 = int32_scalar 0 in\n let v2 = int32_scalar 2 in\n is_true ~msg:\"length 1\" (Space.contains s [ v1 ]);\n is_true ~msg:\"length 3\" (Space.contains s [ v1; v2; v1 ]);\n is_false ~msg:\"empty\" (Space.contains s []);\n is_false ~msg:\"too long\" (Space.contains s [ v1; v2; v1; v2 ])\n\nlet test_seq_contains_unbounded () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Sequence.create ~min_length:0 ds in\n is_true ~msg:\"empty is valid\" (Space.contains s []);\n is_true ~msg:\"long is valid\"\n (Space.contains s (List.init 100 (fun _ -> int32_scalar 0)))\n\nlet test_seq_sample () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Sequence.create ~min_length:1 ~max_length:5 ds in\n let v = Space.sample s in\n is_true ~msg:\"sample valid\" (Space.contains s v)\n\nlet test_seq_sample_fixed () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Sequence.create ~min_length:2 ds in\n let v = Space.sample s in\n equal ~msg:\"fixed length 2\" int 2 (List.length v)\n\nlet test_seq_pack_unpack () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Sequence.create ~min_length:1 ~max_length:3 ds in\n let v = [ int32_scalar 0; int32_scalar 1 ] in\n let packed = Space.pack s v in\n let unpacked = Space.unpack s packed in\n is_ok ~msg:\"round-trip succeeds\" unpacked\n\nlet test_seq_error_min_negative () =\n let ds = Space.Discrete.create 3 in\n raises_invalid_arg \"Space.Sequence.create: min_length must be non-negative\"\n (fun () -> Space.Sequence.create ~min_length:(-1) ds)\n\nlet test_seq_error_max_lt_min () =\n let ds = Space.Discrete.create 3 in\n raises_invalid_arg \"Space.Sequence.create: max_length must be >= min_length\"\n (fun () -> Space.Sequence.create ~min_length:5 ~max_length:2 ds)\n\n(* Discrete helpers *)\n\nlet test_discrete_to_int () =\n let v = Space.Discrete.of_int 5 in\n equal ~msg:\"to_int round-trip\" int 5 (Space.Discrete.to_int v)\n\nlet test_discrete_of_int () =\n let v = Space.Discrete.of_int 3 in\n let s = Space.Discrete.create 5 in\n is_true ~msg:\"of_int creates valid element\" (Space.contains s v)\n\n(* Spec *)\n\nlet test_spec_discrete () =\n let s = Space.Discrete.create ~start:2 4 in\n let sp = Space.spec s in\n equal ~msg:\"discrete spec\" bool true\n (Space.equal_spec sp (Space.Discrete { start = 2; n = 4 }))\n\nlet test_spec_box () =\n let s = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n let sp = Space.spec s in\n equal ~msg:\"box spec\" bool true\n (Space.equal_spec sp (Space.Box { low = [| 0.0 |]; high = [| 1.0 |] }))\n\nlet test_spec_equal_same () =\n let s1 = Space.Discrete.create 3 in\n let s2 = Space.Discrete.create 3 in\n is_true ~msg:\"same spaces equal spec\"\n (Space.equal_spec (Space.spec s1) (Space.spec s2))\n\nlet test_spec_not_equal_different () =\n let s1 = Space.Discrete.create 3 in\n let s2 = Space.Discrete.create 4 in\n is_false ~msg:\"different spaces not equal spec\"\n (Space.equal_spec (Space.spec s1) (Space.spec s2))\n\nlet test_spec_not_equal_kinds () =\n let s1 = Space.Discrete.create 3 in\n let s2 = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n is_false ~msg:\"different kinds not equal spec\"\n (Space.equal_spec (Space.spec s1) (Space.spec s2))\n\nlet test_spec_tuple () =\n let ds = Space.Discrete.create 3 in\n let bs = Space.Box.create ~low:[| 0.0 |] ~high:[| 1.0 |] in\n let s = Space.Tuple.create [ Pack ds; Pack bs ] in\n let sp = Space.spec s in\n let expected =\n Space.Tuple\n [\n Space.Discrete { start = 0; n = 3 };\n Space.Box { low = [| 0.0 |]; high = [| 1.0 |] };\n ]\n in\n is_true ~msg:\"tuple spec\" (Space.equal_spec sp expected)\n\nlet test_spec_dict () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Dict.create [ (\"a\", Pack ds) ] in\n let sp = Space.spec s in\n let expected = Space.Dict [ (\"a\", Space.Discrete { start = 0; n = 3 }) ] in\n is_true ~msg:\"dict spec\" (Space.equal_spec sp expected)\n\n(* Tuple.unpack validation *)\n\nlet test_tuple_unpack_validates_elements () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Tuple.create [ Pack ds ] in\n (* Value.Int 5 is out of range for Discrete(n=3, start=0) *)\n let bad = Value.List [ Value.Int 5 ] in\n is_error ~msg:\"unpack rejects invalid element\" (Space.unpack s bad)\n\nlet test_tuple_unpack_valid () =\n let ds = Space.Discrete.create 3 in\n let s = Space.Tuple.create [ Pack ds ] in\n let good = Value.List [ Value.Int 1 ] in\n is_ok ~msg:\"unpack accepts valid element\" (Space.unpack s good)\n\n(* Entry point *)\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu.Space\"\n [\n group \"Discrete\"\n [\n test \"default start\" test_discrete_default;\n test \"custom start\" test_discrete_custom_start;\n test \"contains valid/invalid\" test_discrete_contains;\n test \"contains with start\" test_discrete_contains_with_start;\n test \"sample\" test_discrete_sample;\n test \"pack/unpack\" test_discrete_pack_unpack;\n test \"unpack invalid\" test_discrete_unpack_invalid;\n test \"boundary values\" test_discrete_boundary_values;\n test \"boundary single\" test_discrete_boundary_single;\n test \"shape\" test_discrete_shape;\n test \"error n=0\" test_discrete_error_zero;\n test \"error n<0\" test_discrete_error_negative;\n test \"to_int round-trip\" test_discrete_to_int;\n test \"of_int valid\" test_discrete_of_int;\n ];\n group \"Box\"\n [\n test \"1d create and bounds\" test_box_1d;\n test \"contains\" test_box_contains;\n test \"sample\" test_box_sample;\n test \"pack/unpack\" test_box_pack_unpack;\n test \"boundary values\" test_box_boundary_values;\n test \"boundary identical\" test_box_boundary_identical;\n test \"shape 1d\" test_box_shape_1d;\n test \"2d\" test_box_2d;\n test \"error empty\" test_box_error_empty;\n test \"error mismatched lengths\" test_box_error_mismatch;\n test \"error low > high\" test_box_error_low_gt_high;\n ];\n group \"Multi_binary\"\n [\n test \"contains\" test_mb_contains;\n test \"sample\" test_mb_sample;\n test \"boundary values\" test_mb_boundary_values;\n test \"shape\" test_mb_shape;\n test \"error n=0\" test_mb_error;\n ];\n group \"Multi_discrete\"\n [\n test \"contains\" test_md_contains;\n test \"sample\" test_md_sample;\n test \"shape\" test_md_shape;\n test \"error empty\" test_md_error_empty;\n test \"error element <= 0\" test_md_error_zero_element;\n ];\n group \"Tuple\"\n [\n test \"contains\" test_tuple_contains;\n test \"sample\" test_tuple_sample;\n test \"pack/unpack\" test_tuple_pack_unpack;\n test \"empty tuple\" test_tuple_empty;\n test \"unpack validates elements\" test_tuple_unpack_validates_elements;\n test \"unpack valid\" test_tuple_unpack_valid;\n ];\n group \"Dict\"\n [\n test \"contains\" test_dict_contains;\n test \"sample\" test_dict_sample;\n test \"error duplicate keys\" test_dict_error_duplicate;\n ];\n group \"Text\"\n [\n test \"contains\" test_text_contains;\n test \"contains invalid charset\" test_text_contains_invalid;\n test \"contains too long\" test_text_contains_too_long;\n test \"sample\" test_text_sample;\n test \"boundary values\" test_text_boundary_values;\n test \"error max_length=0\" test_text_error_max_length;\n test \"error empty charset\" test_text_error_charset;\n ];\n group \"Sequence\"\n [\n test \"contains bounded\" test_seq_contains;\n test \"contains unbounded\" test_seq_contains_unbounded;\n test \"sample\" test_seq_sample;\n test \"sample fixed length\" test_seq_sample_fixed;\n test \"pack/unpack\" test_seq_pack_unpack;\n test \"error min < 0\" test_seq_error_min_negative;\n test \"error max < min\" test_seq_error_max_lt_min;\n ];\n group \"spec\"\n [\n test \"discrete\" test_spec_discrete;\n test \"box\" test_spec_box;\n test \"equal same\" test_spec_equal_same;\n test \"not equal different\" test_spec_not_equal_different;\n test \"not equal kinds\" test_spec_not_equal_kinds;\n test \"tuple\" test_spec_tuple;\n test \"dict\" test_spec_dict;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_value.ml",
"content": "open Fehu\nopen Windtrap\n\nlet value = testable ~pp:Value.pp ~equal:Value.equal ()\n\n(* Equality *)\n\nlet test_null_equal () = equal ~msg:\"null = null\" value Value.Null Value.Null\n\nlet test_bool_equal () =\n equal ~msg:\"true = true\" value (Bool true) (Bool true);\n not_equal ~msg:\"true <> false\" value (Bool true) (Bool false)\n\nlet test_int_equal () =\n equal ~msg:\"1 = 1\" value (Int 1) (Int 1);\n not_equal ~msg:\"1 <> 2\" value (Int 1) (Int 2)\n\nlet test_float_equal () = equal ~msg:\"1.0 = 1.0\" value (Float 1.0) (Float 1.0)\n\nlet test_string_equal () =\n equal ~msg:\"a = a\" value (String \"a\") (String \"a\");\n not_equal ~msg:\"a <> b\" value (String \"a\") (String \"b\")\n\nlet test_int_array_equal () =\n equal ~msg:\"[|1;2|] = [|1;2|]\" value\n (Int_array [| 1; 2 |])\n (Int_array [| 1; 2 |])\n\nlet test_float_array_equal () =\n equal ~msg:\"[|1.0|] = [|1.0|]\" value (Float_array [| 1.0 |])\n (Float_array [| 1.0 |])\n\nlet test_bool_array_equal () =\n equal ~msg:\"[|true|] = [|true|]\" value (Bool_array [| true |])\n (Bool_array [| true |])\n\nlet test_list_equal () =\n equal ~msg:\"[Int 1] = [Int 1]\" value (List [ Int 1 ]) (List [ Int 1 ])\n\nlet test_dict_equal () =\n equal ~msg:\"dict equal\" value (Dict [ (\"k\", Int 1) ]) (Dict [ (\"k\", Int 1) ])\n\nlet test_cross_type_inequality () =\n not_equal ~msg:\"Int 1 <> Float 1.0\" value (Int 1) (Float 1.0);\n not_equal ~msg:\"Null <> Int 0\" value Null (Int 0)\n\n(* Formatting *)\n\nlet test_to_string_null () =\n equal ~msg:\"null\" string \"null\" (Value.to_string Null)\n\nlet test_to_string_bool () =\n equal ~msg:\"bool true\" string \"true\" (Value.to_string (Bool true))\n\nlet test_to_string_int () =\n equal ~msg:\"int 42\" string \"42\" (Value.to_string (Int 42))\n\nlet test_to_string_float () =\n let s = Value.to_string (Float 3.14) in\n is_true ~msg:\"float non-empty\" (String.length s > 0)\n\nlet test_to_string_string () =\n let s = Value.to_string (String \"hello\") in\n is_true ~msg:\"string non-empty\" (String.length s > 0)\n\nlet test_to_string_arrays () =\n is_true ~msg:\"int_array non-empty\"\n (String.length (Value.to_string (Int_array [| 1 |])) > 0);\n is_true ~msg:\"float_array non-empty\"\n (String.length (Value.to_string (Float_array [| 1.0 |])) > 0);\n is_true ~msg:\"bool_array non-empty\"\n (String.length (Value.to_string (Bool_array [| true |])) > 0)\n\nlet test_to_string_list () =\n let s = Value.to_string (List [ Int 1; Int 2 ]) in\n is_true ~msg:\"list non-empty\" (String.length s > 0)\n\nlet test_to_string_dict () =\n let s = Value.to_string (Dict [ (\"k\", Int 1) ]) in\n is_true ~msg:\"dict non-empty\" (String.length s > 0)\n\nlet () =\n run \"Fehu.Value\"\n [\n group \"equality\"\n [\n test \"null\" test_null_equal;\n test \"bool\" test_bool_equal;\n test \"int\" test_int_equal;\n test \"float\" test_float_equal;\n test \"string\" test_string_equal;\n test \"int_array\" test_int_array_equal;\n test \"float_array\" test_float_array_equal;\n test \"bool_array\" test_bool_array_equal;\n test \"list\" test_list_equal;\n test \"dict\" test_dict_equal;\n test \"cross-type inequality\" test_cross_type_inequality;\n ];\n group \"formatting\"\n [\n test \"to_string null\" test_to_string_null;\n test \"to_string bool\" test_to_string_bool;\n test \"to_string int\" test_to_string_int;\n test \"to_string float\" test_to_string_float;\n test \"to_string string\" test_to_string_string;\n test \"to_string arrays\" test_to_string_arrays;\n test \"to_string list\" test_to_string_list;\n test \"to_string dict\" test_to_string_dict;\n ];\n ]\n"
},
{
"path": "packages/fehu/test/test_vec_env.ml",
"content": "open Fehu\nopen Windtrap\n\nlet make_test_env ?(max_steps = 100) () =\n let obs_space = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act_space = Space.Discrete.create 2 in\n let state = ref 5.0 in\n let steps = ref 0 in\n let reset _env ?options:_ () =\n state := 5.0;\n steps := 0;\n (Nx.create Nx.float32 [| 1 |] [| !state |], Info.empty)\n in\n let step _env action =\n let a : Int32.t array = Nx.to_array (Nx.reshape [| 1 |] action) in\n state := !state +. if Int32.to_int a.(0) = 0 then -1.0 else 1.0;\n incr steps;\n let terminated = !state <= 0.0 || !state >= 10.0 in\n let truncated = (not terminated) && !steps >= max_steps in\n Env.step_result\n ~observation:(Nx.create Nx.float32 [| 1 |] [| !state |])\n ~reward:1.0 ~terminated ~truncated ()\n in\n Env.create ~id:\"Test-v0\" ~observation_space:obs_space ~action_space:act_space\n ~reset ~step ()\n\nlet make_envs n = List.init n (fun _ -> make_test_env ())\n\n(* Creation *)\n\nlet test_create_num_envs () =\n let venv = Vec_env.create (make_envs 3) in\n equal ~msg:\"num_envs\" int 3 (Vec_env.num_envs venv)\n\nlet test_spaces_match_first_env () =\n let envs = make_envs 3 in\n let venv = Vec_env.create envs in\n let obs_shape = Space.shape (Vec_env.observation_space venv) in\n let act_shape = Space.shape (Vec_env.action_space venv) in\n let first_obs = Space.shape (Env.observation_space (List.hd envs)) in\n let first_act = Space.shape (Env.action_space (List.hd envs)) in\n equal ~msg:\"obs space shape\" (option (array int)) first_obs obs_shape;\n equal ~msg:\"act space shape\" (option (array int)) first_act act_shape\n\nlet test_empty_list_raises () =\n raises_invalid_arg \"Vec_env.create: env list must not be empty\" (fun () ->\n ignore (Vec_env.create []))\n\nlet test_incompatible_spaces_raises () =\n let obs1 = Space.Box.create ~low:[| 0.0 |] ~high:[| 10.0 |] in\n let act = Space.Discrete.create 2 in\n let obs2 = Space.Box.create ~low:[| 0.0 |] ~high:[| 5.0 |] in\n let make_env obs =\n let reset _env ?options:_ () =\n (Nx.create Nx.float32 [| 1 |] [| 0.0 |], Info.empty)\n in\n let step _env _action =\n Env.step_result ~observation:(Nx.create Nx.float32 [| 1 |] [| 0.0 |]) ()\n in\n Env.create ~observation_space:obs ~action_space:act ~reset ~step ()\n in\n let e1 = make_env obs1 in\n let e2 = make_env obs2 in\n raises_match ~msg:\"incompatible spaces raises\"\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Vec_env.create [ e1; e2 ]))\n\n(* Reset *)\n\nlet test_reset_obs_length () =\n let venv = Vec_env.create (make_envs 3) in\n let obs, _infos = Vec_env.reset venv () in\n equal ~msg:\"obs array length\" int 3 (Array.length obs)\n\nlet test_reset_infos_length () =\n let venv = Vec_env.create (make_envs 3) in\n let _obs, infos = Vec_env.reset venv () in\n equal ~msg:\"infos array length\" int 3 (Array.length infos)\n\n(* Step *)\n\nlet test_step_result_lengths () =\n let venv = Vec_env.create (make_envs 3) in\n let _obs, _infos = Vec_env.reset venv () in\n let action = Nx.create Nx.int32 [| 1 |] [| 1l |] in\n let actions = Array.make 3 action in\n let s = Vec_env.step venv actions in\n equal ~msg:\"observations length\" int 3 (Array.length s.observations);\n equal ~msg:\"rewards length\" int 3 (Array.length s.rewards);\n equal ~msg:\"terminated length\" int 3 (Array.length s.terminated);\n equal ~msg:\"truncated length\" int 3 (Array.length s.truncated);\n equal ~msg:\"infos length\" int 3 (Array.length s.infos)\n\nlet test_wrong_action_count_raises () =\n let venv = Vec_env.create (make_envs 3) in\n let _obs, _infos = Vec_env.reset venv () in\n let action = Nx.create Nx.int32 [| 1 |] [| 1l |] in\n let actions = Array.make 2 action in\n raises_invalid_arg \"Vec_env.step: expected 3 actions, got 2\" (fun () ->\n ignore (Vec_env.step venv actions))\n\nlet test_autoreset_final_observation () =\n let env = make_test_env ~max_steps:3 () in\n let venv = Vec_env.create [ env ] in\n let _obs, _infos = Vec_env.reset venv () in\n let right = Nx.create Nx.int32 [| 1 |] [| 1l |] in\n let actions = [| right |] in\n (* Step until truncated at max_steps=3 *)\n let s1 = Vec_env.step venv actions in\n is_false ~msg:\"not done after step 1\" s1.truncated.(0);\n let s2 = Vec_env.step venv actions in\n is_false ~msg:\"not done after step 2\" s2.truncated.(0);\n let s3 = Vec_env.step venv actions in\n is_true ~msg:\"truncated after step 3\" s3.truncated.(0);\n (* After autoreset, info should have final_observation *)\n is_some ~msg:\"final_observation key present\"\n (Info.find \"final_observation\" s3.infos.(0));\n (* Observation should be from reset (5.0), not terminal *)\n let arr : float array =\n Nx.to_array (Nx.reshape [| 1 |] s3.observations.(0))\n in\n equal ~msg:\"obs is from reset\" (float 1e-6) 5.0 arr.(0)\n\n(* Close *)\n\nlet test_close_all_envs () =\n let envs = make_envs 3 in\n let venv = Vec_env.create envs in\n Vec_env.close venv;\n List.iter (fun env -> is_true ~msg:\"env is closed\" (Env.closed env)) envs\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n run \"Fehu.Vec_env\"\n [\n group \"creation\"\n [\n test \"num_envs\" test_create_num_envs;\n test \"spaces match first env\" test_spaces_match_first_env;\n test \"empty list raises\" test_empty_list_raises;\n test \"incompatible spaces raises\" test_incompatible_spaces_raises;\n ];\n group \"reset\"\n [\n test \"observations length\" test_reset_obs_length;\n test \"infos length\" test_reset_infos_length;\n ];\n group \"step\"\n [\n test \"result array lengths\" test_step_result_lengths;\n test \"wrong action count raises\" test_wrong_action_count_raises;\n test \"autoreset with final_observation\"\n test_autoreset_final_observation;\n ];\n group \"close\" [ test \"closes all inner envs\" test_close_all_envs ];\n ]\n"
},
{
"path": "packages/hugin/README.md",
"content": "# Hugin\n\nDeclarative plotting and visualization library for OCaml.\n\nHugin is part of the Raven ecosystem, providing a functional API to create\npublication-quality charts and figures from Nx arrays. You build immutable\nplot specifications with mark constructors, compose them with `|>` pipelines,\nand render to PNG, SVG, PDF, or an interactive SDL window.\n\n## Features\n\n- Line, scatter, bar, histogram, error bar, fill-between, hline/vline, hspan/vspan\n- Heatmap, colormapped image display (`imshow`), contour plots\n- Multi-panel layouts with `Layout.grid`, `Layout.hstack`, `Layout.vstack`\n- Perceptually uniform OKLCH colors with colorblind-friendly Okabe-Ito palette\n- Predefined colormaps: viridis, plasma, inferno, magma, cividis, coolwarm\n- Themes with context scaling (paper, notebook, talk, poster)\n- Axis scales: linear, log, sqrt, asinh, symlog\n- Cairo rendering (PNG, PDF), pure-OCaml SVG backend, interactive SDL display\n- Format printer for Quill notebooks (`#install_printer`)\n\n## Quick Start\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. 6.28 100 in\n let y = Nx.sin x in\n line ~x ~y () |> title \"Sine wave\" |> render_png \"sine.png\"\n```\n\n## Contributing\n\nSee the [Raven monorepo README](../../README.md) for guidelines.\n\n## License\n\nISC License. See [LICENSE](../../LICENSE) for details.\n"
},
{
"path": "packages/hugin/doc/01-getting-started.md",
"content": "# Getting Started\n\nThis guide covers installation, your first plot, and the key concepts behind Hugin.\n\n## Installation\n\nInstall system dependencies:\n\n\n```bash\n# macOS\nbrew install cairo sdl2\n\n# Ubuntu/Debian\napt install libcairo2-dev libsdl2-dev\n```\n\nThen install hugin:\n\n\n```bash\nopam install hugin\n```\n\nOr build from source:\n\n\n```bash\ngit clone https://github.com/raven-ml/raven\ncd raven && dune build dev/hugin\n```\n\nAdd to your `dune` file:\n\n\n```dune\n(executable\n (name main)\n (libraries hugin))\n```\n\n## Your First Plot\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 100 in\n let y = Nx.sin x in\n line ~x ~y () |> title \"Sine wave\" |> render_png \"sine.png\"\n```\n\nThis creates a 1-D array of 100 points, computes the sine, builds a line specification, adds a title, and writes a PNG file.\n\n## Key Concepts\n\n### Marks\n\nA mark constructor (`line`, `point`, `bar`, `hist`, `heatmap`, etc.) takes data arrays and optional visual properties and returns an immutable plot specification of type `t`. A mark is already a complete spec — you can render it directly:\n\n\n```ocaml\nline ~x ~y () |> render_png \"plot.png\"\n```\n\n### Decorations\n\nDecoration functions add metadata to a spec. They are designed for the `|>` pipeline:\n\n\n```ocaml\nline ~x ~y ()\n|> title \"My Plot\"\n|> xlabel \"Time (s)\"\n|> ylabel \"Amplitude\"\n|> xlim 0. 10.\n|> grid_lines true\n```\n\nDecorations include `title`, `xlabel`, `ylabel`, `xlim`, `ylim`, `xscale`, `yscale`, `grid_lines`, `legend`, `xticks`, `yticks`, `xinvert`, `yinvert`, `with_theme`, and tick formatting.\n\n### Composition\n\n`layers` overlays multiple marks on shared axes:\n\n\n```ocaml\nlayers [\n line ~x ~y:(Nx.sin x) ~label:\"sin\" ();\n line ~x ~y:(Nx.cos x) ~label:\"cos\" ~line_style:`Dashed ();\n]\n|> legend |> render_png \"overlay.png\"\n```\n\nYou can mix mark types freely. A `line` with `point` markers, a `bar` chart with `hline` reference lines — anything goes.\n\n### Layout\n\n`Layout.grid` arranges specs in rows and columns:\n\n\n```ocaml\nlet p1 = line ~x ~y:(Nx.sin x) () |> title \"sin\" in\nlet p2 = line ~x ~y:(Nx.cos x) () |> title \"cos\" in\nLayout.grid [ [ p1; p2 ] ] |> render_png \"grid.png\"\n```\n\n`Layout.hstack` and `Layout.vstack` are shorthands for single-row and single-column grids.\n\n### Rendering\n\nFour output modes:\n\n| Function | Output |\n|----------|--------|\n| `render_png \"file.png\" t` | PNG image file |\n| `render_svg \"file.svg\" t` | SVG document file |\n| `render_pdf \"file.pdf\" t` | PDF document file |\n| `show t` | Interactive SDL window (resize, Esc to close) |\n\nAll renderers accept optional `~width` and `~height` (default 1600×1200) and `~theme`.\n\n`render_svg_to_string` and `render_to_buffer` return the output as a string instead of writing a file.\n\n## Common Marks\n\n### Line\n\n\n```ocaml\nline ~x ~y ()\nline ~x ~y ~color:Color.blue ~line_style:`Dashed ~line_width:2.0 ()\nline ~x ~y ~step:`Post () (* staircase plot *)\n```\n\n### Scatter\n\n\n```ocaml\npoint ~x ~y ()\npoint ~x ~y ~color_by:values ~size:8. ~marker:Star ()\npoint ~x ~y ~size_by:weights () (* variable marker size *)\n```\n\n### Bar Chart\n\n\n```ocaml\nbar ~x:categories ~height:values ()\nbar ~x:categories ~height:values ~width:0.5 ~color:Color.orange ()\n```\n\n### Histogram\n\n\n```ocaml\nhist ~x:data ()\nhist ~x:data ~bins:(`Num 30) ~density:true ~color:Color.green ()\n```\n\n### Heatmap\n\n\n```ocaml\n(* data has shape [|rows; cols|] *)\nheatmap ~data ()\nheatmap ~data ~annotate:true ~cmap:Cmap.viridis ()\n```\n\n### Fill Between\n\n\n```ocaml\nfill_between ~x ~y1:(Nx.sub y err) ~y2:(Nx.add y err) ~alpha:0.3 ()\n```\n\n### Error Bars\n\n\n```ocaml\nerrorbar ~x ~y ~yerr:(`Symmetric err) ()\nerrorbar ~x ~y ~yerr:(`Asymmetric (lo, hi)) ~xerr:(`Symmetric xerr) ()\n```\n\n## Next Steps\n\n- [Marks and Styling](/docs/hugin/marks-and-styling/) — full mark catalog and visual properties\n- [Layout and Decorations](/docs/hugin/layout-and-decorations/) — axes, scales, themes, multi-panel\n- [Colors and Colormaps](/docs/hugin/colors-and-colormaps/) — OKLCH colors and colormap reference\n"
},
{
"path": "packages/hugin/doc/02-marks-and-styling.md",
"content": "# Marks and Styling\n\nEvery visualization in Hugin starts with one or more marks. A mark constructor takes data arrays and optional visual properties and returns an immutable plot specification.\n\n## Mark Catalog\n\n### Line Plots\n\n`line ~x ~y ()` connects points `(x.(i), y.(i))` with straight segments.\n\n| Argument | Type | Default | Description |\n|----------|------|---------|-------------|\n| `~x` | `Nx.float32_t` | required | X coordinates |\n| `~y` | `Nx.float32_t` | required | Y coordinates |\n| `~color` | `Color.t` | theme palette | Line color |\n| `~line_width` | `float` | theme line width | Stroke width |\n| `~line_style` | `` `Solid \\| `Dashed \\| `Dotted \\| `Dash_dot `` | `` `Solid `` | Dash pattern |\n| `~step` | `` `Pre \\| `Post \\| `Mid `` | none | Staircase interpolation |\n| `~marker` | `marker` | none | Marker at each point |\n| `~label` | `string` | none | Legend entry |\n| `~alpha` | `float` | 1.0 | Opacity |\n\nStep modes: `Post` holds each value until the next x-point, `Pre` steps at the current x-point, `Mid` steps at the midpoint.\n\n### Scatter Plots\n\n`point ~x ~y ()` places individual markers at data coordinates.\n\n| Argument | Type | Default | Description |\n|----------|------|---------|-------------|\n| `~x` | `Nx.float32_t` | required | X coordinates |\n| `~y` | `Nx.float32_t` | required | Y coordinates |\n| `~color` | `Color.t` | theme palette | Uniform color |\n| `~color_by` | `Nx.float32_t` | none | Per-point values mapped through sequential colormap |\n| `~size` | `float` | theme marker size | Uniform marker size |\n| `~size_by` | `Nx.float32_t` | none | Per-point values for variable marker area |\n| `~marker` | `marker` | `Circle` | Marker shape |\n| `~label` | `string` | none | Legend entry |\n| `~alpha` | `float` | 1.0 | Opacity |\n\nWhen `~color_by` is set, a colorbar is displayed showing the value-to-color mapping.\n\n### Bar Charts\n\n`bar ~x ~height ()` draws vertical bars centered on `x` values.\n\n| Argument | Type | Default | Description |\n|----------|------|---------|-------------|\n| `~x` | `Nx.float32_t` | required | Bar center positions |\n| `~height` | `Nx.float32_t` | required | Bar heights |\n| `~width` | `float` | 0.8 | Bar width |\n| `~bottom` | `float` | 0.0 | Baseline y-value |\n| `~color` | `Color.t` | theme palette | Fill color |\n| `~label` | `string` | none | Legend entry |\n| `~alpha` | `float` | 1.0 | Opacity |\n\n### Histograms\n\n`hist ~x ()` bins the values in `x` and draws a bar chart.\n\n| Argument | Type | Default | Description |\n|----------|------|---------|-------------|\n| `~x` | `Nx.float32_t` | required | Data values |\n| `~bins` | `` `Num of int \\| `Edges of float array `` | `` `Num 10 `` | Number of bins or explicit edges |\n| `~density` | `bool` | false | Normalize so total area equals 1.0 |\n| `~color` | `Color.t` | theme palette | Fill color |\n| `~label` | `string` | none | Legend entry |\n\n### Reference Lines and Spans\n\n`hline ~y ()` draws a horizontal line across the full plot width. `vline ~x ()` draws a vertical line across the full height. Both accept `~color`, `~line_width`, `~line_style`, `~label`, and `~alpha`.\n\n`hspan ~y0 ~y1 ()` shades a horizontal band. `vspan ~x0 ~x1 ()` shades a vertical band. Both accept `~color`, `~alpha` (default 0.2), and `~label`.\n\n### Fill Between\n\n`fill_between ~x ~y1 ~y2 ()` fills the area between two curves. `~alpha` defaults to 0.3.\n\n### Error Bars\n\n`errorbar ~x ~y ~yerr ()` draws error bars at each point.\n\n- `~yerr`: `` `Symmetric e `` draws y ± e, `` `Asymmetric (lo, hi) `` draws [y - lo, y + hi]\n- `~xerr`: optional horizontal error bars with the same format\n- `~cap_size`: cap width (defaults to half the theme marker size)\n\n### Text\n\n`text ~x ~y \"label\" ()` places a string at data coordinates `(x, y)`. Accepts `~color` and `~font_size`.\n\n### Image\n\n`image data` displays an Nx uint8 array as an image. `data` must have shape `[|h; w; 3|]` (RGB) or `[|h; w; 4|]` (RGBA).\n\n### Colormapped Image\n\n`imshow ~data ()` displays a 2-D float array through a colormap.\n\n| Argument | Type | Default | Description |\n|----------|------|---------|-------------|\n| `~data` | `Nx.float32_t` | required | 2-D array of shape `[|rows; cols|]` |\n| `~stretch` | `` `Linear \\| `Log \\| `Sqrt \\| `Asinh \\| `Power of float `` | `` `Linear `` | Transfer function before colormap lookup |\n| `~cmap` | `Cmap.t` | theme sequential | Colormap |\n| `~vmin` | `float` | data min | Lower bound of color range |\n| `~vmax` | `float` | data max | Upper bound of color range |\n\n### Heatmap\n\n`heatmap ~data ()` displays a 2-D array as a grid of colored cells. Row 0 appears at the top.\n\n| Argument | Type | Default | Description |\n|----------|------|---------|-------------|\n| `~data` | `Nx.float32_t` | required | 2-D array of shape `[|rows; cols|]` |\n| `~annotate` | `bool` | false | Show cell values |\n| `~cmap` | `Cmap.t` | theme sequential | Colormap |\n| `~vmin` | `float` | data min | Lower bound |\n| `~vmax` | `float` | data max | Upper bound |\n| `~fmt` | `float -> string` | `Printf.sprintf \"%.2g\"` | Cell value formatter (when annotate is true) |\n\n### Contour\n\n`contour ~data ~x0 ~x1 ~y0 ~y1 ()` draws iso-level contour lines through a 2-D grid.\n\n| Argument | Type | Default | Description |\n|----------|------|---------|-------------|\n| `~data` | `Nx.float32_t` | required | 2-D array of shape `[|rows; cols|]` |\n| `~x0`, `~x1`, `~y0`, `~y1` | `float` | required | Data-space rectangle |\n| `~levels` | `` `Num of int \\| `Values of float array `` | `` `Num 8 `` | Number of levels or explicit values |\n| `~filled` | `bool` | false | Fill regions between levels |\n| `~cmap` | `Cmap.t` | theme sequential | Per-level colormap |\n| `~color` | `Color.t` | none | Single stroke color (unfilled contours) |\n| `~line_width` | `float` | theme line width | Stroke width |\n| `~label` | `string` | none | Legend entry |\n| `~alpha` | `float` | 1.0 | Opacity |\n\n## Marker Shapes\n\nFive built-in shapes:\n\n| Marker | Description |\n|--------|-------------|\n| `Circle` | Filled circle |\n| `Square` | Filled square |\n| `Triangle` | Filled triangle |\n| `Plus` | Plus sign (+) |\n| `Star` | Five-pointed star |\n\nUse with `line ~marker:Triangle` or `point ~marker:Star`.\n\n## Auto-Coloring\n\nWhen you omit `~color`, marks are colored automatically from the theme's categorical palette. The first mark in a spec gets `palette.(0)`, the second gets `palette.(1)`, and so on. Explicitly setting `~color` takes precedence.\n\n## Next Steps\n\n- [Layout and Decorations](/docs/hugin/layout-and-decorations/) — axes, scales, themes, multi-panel layouts\n- [Colors and Colormaps](/docs/hugin/colors-and-colormaps/) — OKLCH color space, palettes, and colormap reference\n"
},
{
"path": "packages/hugin/doc/03-layout-and-decorations.md",
"content": "# Layout and Decorations\n\nDecorations add metadata and styling to a plot specification. Layout functions arrange multiple specs into multi-panel figures.\n\n## Decorations\n\nAll decoration functions take a `t` and return a new `t`, designed for the `|>` pipeline:\n\n\n```ocaml\nline ~x ~y ()\n|> title \"Frequency Response\"\n|> xlabel \"Frequency (Hz)\"\n|> ylabel \"Magnitude (dB)\"\n|> xscale `Log\n|> ylim (-60.) 0.\n|> grid_lines true\n```\n\n### Titles and Labels\n\n| Function | Description |\n|----------|-------------|\n| `title s t` | Plot title |\n| `xlabel s t` | X-axis label |\n| `ylabel s t` | Y-axis label |\n\n### Axis Limits\n\n| Function | Description |\n|----------|-------------|\n| `xlim lo hi t` | Fix x-axis range to [lo, hi] |\n| `ylim lo hi t` | Fix y-axis range to [lo, hi] |\n\nWhen omitted, axis ranges are computed automatically from the data with 5% padding.\n\n### Axis Scales\n\n| Function | Description |\n|----------|-------------|\n| `xscale s t` | Set x-axis scale |\n| `yscale s t` | Set y-axis scale |\n\nAvailable scales:\n\n| Scale | When to use |\n|-------|-------------|\n| `` `Linear `` | Default. Uniform spacing. |\n| `` `Log `` | Data spanning multiple orders of magnitude. All values must be positive. |\n| `` `Sqrt `` | Moderate compression of large values. Handles zero. |\n| `` `Asinh `` | Like log but handles zero and negative values. Transitions smoothly from linear near zero to logarithmic at large magnitudes. |\n| `` `Symlog linthresh `` | Linear within [-linthresh, linthresh], logarithmic outside. Good for data with both small and large values centered around zero. |\n\n### Axis Direction\n\n| Function | Description |\n|----------|-------------|\n| `xinvert t` | X-axis values increase right-to-left |\n| `yinvert t` | Y-axis values increase top-to-bottom |\n\nUseful for conventions like right ascension in sky charts (xinvert) or magnitude axes in HR diagrams (yinvert).\n\n### Ticks\n\n| Function | Description |\n|----------|-------------|\n| `xticks ticks t` | Explicit tick positions and labels as `(float * string) list` |\n| `yticks ticks t` | Same for y-axis |\n| `xtick_format fmt t` | Custom tick label formatter (preserves auto-generated positions) |\n| `ytick_format fmt t` | Same for y-axis |\n\nExample with explicit ticks:\n\n\n```ocaml\nline ~x ~y ()\n|> xticks [ (0., \"Jan\"); (1., \"Feb\"); (2., \"Mar\"); (3., \"Apr\") ]\n```\n\nExample with custom formatting:\n\n\n```ocaml\nline ~x ~y ()\n|> xtick_format (Printf.sprintf \"%.1f%%\")\n```\n\n### Grid and Legend\n\n| Function | Description |\n|----------|-------------|\n| `grid_lines visible t` | Show or hide grid lines |\n| `legend ?loc t` | Show legend at `loc` (default `Upper_right`) |\n\nLegend locations: `Upper_right`, `Upper_left`, `Lower_right`, `Lower_left`, `Center`.\n\nThe legend is populated from marks that have a `~label`. Marks without labels are excluded.\n\n### Theme Override\n\n`with_theme theme t` renders with `theme` instead of the default.\n\n## Layout\n\n### Grid\n\n`Layout.grid rows` arranges specs in a grid where each inner list is a row:\n\n\n```ocaml\nlet p1 = line ~x ~y:(Nx.sin x) () |> title \"sin\" in\nlet p2 = line ~x ~y:(Nx.cos x) () |> title \"cos\" in\nlet p3 = line ~x ~y:(Nx.tan x) () |> title \"tan\" |> ylim (-5.) 5. in\nlet p4 = hist ~x:(Nx.rand Nx.float32 [|500|]) () |> title \"random\" in\nLayout.grid [ [ p1; p2 ]; [ p3; p4 ] ] |> render_png \"grid.png\"\n```\n\n`~gap` controls spacing between panels as a fraction of total size (default 0.05).\n\n### Stack\n\n| Function | Description |\n|----------|-------------|\n| `Layout.hstack specs` | Single row of panels |\n| `Layout.vstack specs` | Single column of panels |\n\nBoth accept `~gap`.\n\n## Themes\n\nA theme controls every non-data visual element: background, typography, axes, grid, spacing, and data palettes.\n\n### Predefined Themes\n\n| Theme | Description |\n|-------|-------------|\n| `Theme.default` | Light background, subtle grid, Okabe-Ito palette |\n| `Theme.dark` | Dark background |\n| `Theme.minimal` | No grid, thin axes |\n\n\n```ocaml\nline ~x ~y () |> with_theme Theme.dark |> render_png \"dark.png\"\n```\n\n### Context Scaling\n\nContext functions scale all visual elements (fonts, line widths, spacing) for different output media:\n\n| Function | Scale factor | Use case |\n|----------|-------------|----------|\n| `Theme.paper` | 1.0 | Journal figures |\n| `Theme.notebook` | 1.3 | Quill notebooks |\n| `Theme.talk` | 1.6 | Slides and presentations |\n| `Theme.poster` | 2.0 | Conference posters |\n\n\n```ocaml\nlet theme = Theme.dark |> Theme.talk in\nline ~x ~y () |> with_theme theme |> render_png \"slide.png\"\n```\n\n### Theme Fields\n\nThe `Theme.t` record is fully public. You can create custom themes by modifying fields:\n\n| Field | Type | Description |\n|-------|------|-------------|\n| `background` | `Color.t` | Background color |\n| `palette` | `Color.t array` | Categorical color palette |\n| `sequential` | `Cmap.t` | Default sequential colormap |\n| `diverging` | `Cmap.t` | Default diverging colormap |\n| `font_title` | `Theme.font` | Title font |\n| `font_label` | `Theme.font` | Axis label font |\n| `font_tick` | `Theme.font` | Tick label font |\n| `axis` | `Theme.line` | Axis line style |\n| `grid` | `Theme.line option` | Grid line style (None to hide) |\n| `tick_length` | `float` | Tick mark length |\n| `padding` | `float` | Plot area padding |\n| `title_gap` | `float` | Gap below title |\n| `label_gap` | `float` | Gap between label and axis |\n| `scale_factor` | `float` | Global size multiplier |\n| `line_width` | `float` | Default line width |\n| `marker_size` | `float` | Default marker size |\n\n## Next Steps\n\n- [Colors and Colormaps](/docs/hugin/colors-and-colormaps/) — OKLCH color space, operations, and colormap reference\n- [Matplotlib Comparison](/docs/hugin/matplotlib-comparison/) — side-by-side with Python\n"
},
{
"path": "packages/hugin/doc/04-colors-and-colormaps.md",
"content": "# Colors and Colormaps\n\nHugin uses the OKLCH color space for perceptually uniform color operations and ships with colorblind-friendly palettes and scientific colormaps.\n\n## Colors\n\n### OKLCH Color Space\n\nColors are represented internally in [OKLCH](https://bottosson.github.io/posts/oklab/), a perceptually uniform color space. Operations like `lighten`, `darken`, and `mix` produce visually consistent results: equal numerical steps yield equal perceived differences.\n\nOKLCH components:\n\n| Component | Range | Description |\n|-----------|-------|-------------|\n| Lightness (L) | [0, 1] | Black to white |\n| Chroma (C) | [0, ~0.4] | Gray to saturated |\n| Hue (H) | [0, 360) | Color wheel angle |\n| Alpha (A) | [0, 1] | Transparency |\n\n### Constructors\n\n\n```ocaml\n(* From OKLCH components *)\nColor.oklch ~l:0.7 ~c:0.15 ~h:145. ()\nColor.oklcha ~l:0.7 ~c:0.15 ~h:145. ~a:0.5 ()\n\n(* From sRGB [0, 1] *)\nColor.rgb ~r:0.2 ~g:0.6 ~b:0.8 ()\nColor.rgba ~r:0.2 ~g:0.6 ~b:0.8 ~a:0.5 ()\n\n(* From hex string *)\nColor.hex \"#3399CC\"\nColor.hex \"#3399CCAA\" (* with alpha *)\n```\n\nAll constructors convert to OKLCH on creation. The reverse conversion (`to_rgba`) is called at render time.\n\n### Accessors\n\n\n```ocaml\nColor.lightness c (* OKLCH lightness *)\nColor.chroma c (* OKLCH chroma *)\nColor.hue c (* OKLCH hue in degrees *)\nColor.alpha c (* alpha channel *)\nColor.to_rgba c (* sRGB (r, g, b, a) tuple, clamped to gamut *)\n```\n\n### Operations\n\n\n```ocaml\nColor.lighten 0.1 c (* increase lightness by 0.1, clamped to [0, 1] *)\nColor.darken 0.1 c (* decrease lightness by 0.1, clamped to [0, 1] *)\nColor.with_alpha 0.5 c (* set alpha *)\nColor.mix 0.5 a b (* blend a and b: 0.0 = a, 1.0 = b *)\n```\n\n`mix` interpolates all OKLCH components. Hue follows the shortest arc on the color wheel.\n\n### Named Colors\n\nThe default named colors follow the [Okabe-Ito palette](https://jfly.uni-koeln.de/color/), designed to be distinguishable under all forms of color-vision deficiency:\n\n| Color | Value |\n|-------|-------|\n| `Color.orange` | Okabe-Ito orange |\n| `Color.sky_blue` | Okabe-Ito sky blue |\n| `Color.green` | Okabe-Ito bluish green |\n| `Color.yellow` | Okabe-Ito yellow |\n| `Color.blue` | Okabe-Ito blue |\n| `Color.vermillion` | Okabe-Ito vermillion |\n| `Color.purple` | Okabe-Ito reddish purple |\n| `Color.black` | Black |\n| `Color.white` | White |\n| `Color.gray` | Neutral gray |\n\n### Formatting\n\n`Color.pp` formats as `oklch(L C H / A)` for debugging.\n\n## Colormaps\n\nA colormap is a continuous mapping from [0, 1] to `Color.t`. Internally stored as a 256-entry lookup table with OKLCH interpolation.\n\n### Evaluation\n\n\n```ocaml\nlet c = Cmap.eval Cmap.viridis 0.5 (* color at midpoint *)\n```\n\nValues are clamped to [0, 1].\n\n### Predefined Colormaps\n\nPerceptually uniform sequential colormaps from the [viridis family](https://bids.github.io/colormap/):\n\n| Colormap | Description |\n|----------|-------------|\n| `Cmap.viridis` | Purple-teal-yellow (default) |\n| `Cmap.plasma` | Purple-orange-yellow |\n| `Cmap.inferno` | Black-purple-orange-yellow |\n| `Cmap.magma` | Black-purple-pink-yellow |\n| `Cmap.cividis` | Optimized for color-vision deficiency |\n\nOther colormaps:\n\n| Colormap | Description |\n|----------|-------------|\n| `Cmap.coolwarm` | Blue-white-red diverging |\n| `Cmap.gray` | Black to white |\n| `Cmap.gray_r` | White to black (standard for astronomy) |\n| `Cmap.hot` | Black-red-yellow-white |\n\n### Custom Colormaps\n\n`Cmap.of_colors` creates a colormap by interpolating linearly through an array of color stops in OKLCH space:\n\n\n```ocaml\nlet my_cmap = Cmap.of_colors [|\n Color.hex \"#000080\";\n Color.hex \"#FFFFFF\";\n Color.hex \"#800000\";\n|]\n```\n\nStops are evenly spaced from 0 to 1. Requires at least 2 colors.\n\n## Using Colors with Marks\n\n### Uniform Color\n\nSet `~color` on any mark:\n\n\n```ocaml\nline ~x ~y ~color:Color.vermillion ()\nbar ~x ~height ~color:(Color.hex \"#336699\") ()\n```\n\n### Data-Driven Color\n\n`point` supports `~color_by` to map per-point values through the theme's sequential colormap:\n\n\n```ocaml\npoint ~x ~y ~color_by:temperature ~marker:Circle ()\n```\n\nA colorbar is displayed automatically.\n\n### Colormaps on 2-D Data\n\n`heatmap`, `imshow`, and `contour` accept `~cmap` to override the default:\n\n\n```ocaml\nheatmap ~data ~cmap:Cmap.coolwarm ()\nimshow ~data ~cmap:Cmap.inferno ~stretch:`Log ()\ncontour ~data ~x0 ~x1 ~y0 ~y1 ~filled:true ~cmap:Cmap.plasma ()\n```\n\n## Next Steps\n\n- [Matplotlib Comparison](/docs/hugin/matplotlib-comparison/) — side-by-side with Python\n- [Marks and Styling](/docs/hugin/marks-and-styling/) — full mark catalog\n"
},
{
"path": "packages/hugin/doc/05-matplotlib-comparison.md",
"content": "# Hugin vs Matplotlib\n\nSide-by-side examples comparing Hugin (OCaml) with Matplotlib (Python). Hugin uses a declarative, pipeline-oriented API while Matplotlib uses an imperative, object-oriented approach.\n\n## Key Differences\n\n| | Hugin | Matplotlib |\n|---|---|---|\n| Style | Declarative, immutable specs | Imperative, mutable state |\n| Composition | `\\|>` pipeline | Method calls on axes |\n| State | No global state | `plt` global state |\n| Colors | OKLCH color space | sRGB strings |\n| Output | `render_png`, `render_svg`, `show` | `plt.savefig`, `plt.show` |\n\n## Line Plot\n\n**Hugin:**\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 100 in\n layers [\n line ~x ~y:(Nx.sin x) ~label:\"sin(x)\" ~color:Color.blue ();\n line ~x ~y:(Nx.cos x) ~label:\"cos(x)\" ~color:Color.vermillion\n ~line_style:`Dashed ();\n ]\n |> title \"Trigonometric Functions\"\n |> xlabel \"Angle (radians)\"\n |> ylabel \"Value\"\n |> ylim (-1.2) 1.2\n |> grid_lines true\n |> legend\n |> render_png \"trig.png\"\n```\n\n**Matplotlib:**\n\n```python\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nx = np.linspace(0, 2 * np.pi, 100)\n\nplt.figure()\nplt.plot(x, np.sin(x), label=\"sin(x)\", color=\"blue\")\nplt.plot(x, np.cos(x), label=\"cos(x)\", color=\"red\", linestyle=\"--\")\nplt.title(\"Trigonometric Functions\")\nplt.xlabel(\"Angle (radians)\")\nplt.ylabel(\"Value\")\nplt.ylim(-1.2, 1.2)\nplt.grid(True)\nplt.legend()\nplt.savefig(\"trig.png\")\n```\n\n## Scatter Plot\n\n**Hugin:**\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.rand Nx.float32 [| 200 |] in\n let y = Nx.rand Nx.float32 [| 200 |] in\n let c = Nx.add x y in\n point ~x ~y ~color_by:c ~size:8. ~marker:Circle ()\n |> title \"Random Scatter\"\n |> render_png \"scatter.png\"\n```\n\n**Matplotlib:**\n\n```python\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nx = np.random.rand(200)\ny = np.random.rand(200)\nc = x + y\n\nplt.figure()\nplt.scatter(x, y, c=c, s=64, marker=\"o\")\nplt.title(\"Random Scatter\")\nplt.colorbar()\nplt.savefig(\"scatter.png\")\n```\n\n## Bar Chart\n\n**Hugin:**\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.create Nx.float32 [| 4 |] [| 1.; 2.; 3.; 4. |] in\n let h = Nx.create Nx.float32 [| 4 |] [| 3.; 7.; 2.; 5. |] in\n bar ~x ~height:h ~color:Color.orange ()\n |> title \"Quarterly Revenue\"\n |> xticks [ (1., \"Q1\"); (2., \"Q2\"); (3., \"Q3\"); (4., \"Q4\") ]\n |> ylabel \"Revenue ($M)\"\n |> render_png \"bar.png\"\n```\n\n**Matplotlib:**\n\n```python\nimport matplotlib.pyplot as plt\n\nx = [1, 2, 3, 4]\nh = [3, 7, 2, 5]\n\nplt.figure()\nplt.bar(x, h, color=\"orange\")\nplt.title(\"Quarterly Revenue\")\nplt.xticks(x, [\"Q1\", \"Q2\", \"Q3\", \"Q4\"])\nplt.ylabel(\"Revenue ($M)\")\nplt.savefig(\"bar.png\")\n```\n\n## Histogram\n\n**Hugin:**\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let data = Nx.randn Nx.float32 [| 1000 |] in\n hist ~x:data ~bins:(`Num 30) ~density:true ~color:Color.sky_blue ()\n |> title \"Normal Distribution\"\n |> xlabel \"Value\"\n |> ylabel \"Density\"\n |> render_png \"hist.png\"\n```\n\n**Matplotlib:**\n\n```python\nimport numpy as np\nimport matplotlib.pyplot as plt\n\ndata = np.random.randn(1000)\n\nplt.figure()\nplt.hist(data, bins=30, density=True, color=\"skyblue\")\nplt.title(\"Normal Distribution\")\nplt.xlabel(\"Value\")\nplt.ylabel(\"Density\")\nplt.savefig(\"hist.png\")\n```\n\n## Multi-Panel Layout\n\n**Hugin:**\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 100 in\n let p1 = line ~x ~y:(Nx.sin x) () |> title \"sin\" in\n let p2 = line ~x ~y:(Nx.cos x) () |> title \"cos\" in\n let p3 = line ~x ~y:(Nx.tan x) () |> title \"tan\" |> ylim (-5.) 5. in\n let p4 = hist ~x:(Nx.rand Nx.float32 [| 500 |]) () |> title \"random\" in\n Layout.grid [ [ p1; p2 ]; [ p3; p4 ] ]\n |> render_png \"grid.png\"\n```\n\n**Matplotlib:**\n\n```python\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nx = np.linspace(0, 2 * np.pi, 100)\n\nfig, axes = plt.subplots(2, 2)\naxes[0, 0].plot(x, np.sin(x)); axes[0, 0].set_title(\"sin\")\naxes[0, 1].plot(x, np.cos(x)); axes[0, 1].set_title(\"cos\")\naxes[1, 0].plot(x, np.tan(x)); axes[1, 0].set_title(\"tan\")\naxes[1, 0].set_ylim(-5, 5)\naxes[1, 1].hist(np.random.rand(500)); axes[1, 1].set_title(\"random\")\nplt.tight_layout()\nplt.savefig(\"grid.png\")\n```\n\n## Heatmap\n\n**Hugin:**\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let data = Nx.init Nx.float32 [| 8; 10 |] (fun idx ->\n let i = Float.of_int idx.(0) and j = Float.of_int idx.(1) in\n Float.sin (i *. 0.5) *. Float.cos (j *. 0.4))\n in\n heatmap ~data ~annotate:true ~cmap:Cmap.viridis ()\n |> title \"Heatmap\"\n |> render_png \"heatmap.png\"\n```\n\n**Matplotlib:**\n\n```python\nimport numpy as np\nimport matplotlib.pyplot as plt\n\ndata = np.fromfunction(\n lambda i, j: np.sin(i * 0.5) * np.cos(j * 0.4), (8, 10)\n)\n\nfig, ax = plt.subplots()\nim = ax.imshow(data, cmap=\"viridis\")\nfor i in range(8):\n for j in range(10):\n ax.text(j, i, f\"{data[i, j]:.2g}\", ha=\"center\", va=\"center\")\nax.set_title(\"Heatmap\")\nplt.colorbar(im)\nplt.savefig(\"heatmap.png\")\n```\n\n## Styling\n\n**Hugin:**\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 50 in\n line ~x ~y:(Nx.sin x)\n ~color:Color.vermillion\n ~line_style:`Dashed\n ~line_width:2.5\n ~marker:Triangle\n ~alpha:0.7\n ()\n |> render_png \"styled.png\"\n```\n\n**Matplotlib:**\n\n```python\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nx = np.linspace(0, 2 * np.pi, 50)\n\nplt.figure()\nplt.plot(x, np.sin(x), color=\"red\", linestyle=\"--\",\n linewidth=2.5, marker=\"^\", alpha=0.7)\nplt.savefig(\"styled.png\")\n```\n\n## Themes\n\nHugin provides built-in themes with context scaling. Matplotlib uses style sheets.\n\n**Hugin:**\n\n\n```ocaml\n(* Dark theme scaled for a presentation *)\nlet theme = Theme.dark |> Theme.talk in\nline ~x ~y () |> with_theme theme |> render_png \"slide.png\"\n```\n\n**Matplotlib:**\n\n```python\nplt.style.use(\"dark_background\")\nplt.rcParams.update({\"font.size\": 14})\nplt.plot(x, y)\nplt.savefig(\"slide.png\")\n```\n\n## Save and Export\n\n**Hugin:**\n\n\n```ocaml\nlet spec = line ~x ~y () |> title \"My Plot\" in\nspec |> render_png \"plot.png\";\nspec |> render_svg \"plot.svg\";\nspec |> render_pdf \"plot.pdf\";\nspec |> show (* interactive SDL window *)\n```\n\n**Matplotlib:**\n\n```python\nplt.plot(x, y)\nplt.title(\"My Plot\")\nplt.savefig(\"plot.png\")\nplt.savefig(\"plot.svg\")\nplt.savefig(\"plot.pdf\")\nplt.show()\n```\n\nIn Hugin, the spec is an immutable value. You can render the same spec to multiple formats without rebuilding it. In Matplotlib, the figure is mutable state that `savefig` and `show` consume.\n\n## Interactive Display\n\n**Hugin:**\n\n\n```ocaml\nshow ~width:1600. ~height:1200. spec\n```\n\nThe SDL window is resizable. The plot re-renders at the new dimensions. Press Escape or Q to close.\n\n**Matplotlib:**\n\n```python\nplt.show()\n```\n"
},
{
"path": "packages/hugin/doc/index.md",
"content": "# Hugin\n\nHugin creates publication-quality plots from Nx arrays using a declarative, pipeline-oriented API.\n\n## What Hugin Does\n\nHugin turns immutable plot specifications into rendered output. You build a specification from mark constructors (`line`, `point`, `bar`, `hist`), decorate it with `title`, `xlabel`, and axis controls via the `|>` pipeline, and render with `render_png`, `render_svg`, or `show`.\n\nInternally, rendering proceeds in three stages: the user-facing spec is compiled to a prepared tree (histograms binned, data bounds computed, marks auto-colored), then resolved to device-pixel coordinates, then drawn by a backend. Data compilation happens once; layout resolution is cheap and repeatable at different sizes.\n\n## System Requirements\n\nHugin needs Cairo and SDL2 for rendering:\n\n\n```bash\n# macOS\nbrew install cairo sdl2\n\n# Ubuntu/Debian\napt install libcairo2-dev libsdl2-dev\n```\n\n## Quick Start\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 100 in\n let y = Nx.sin x in\n line ~x ~y () |> title \"Sine wave\" |> render_png \"sine.png\"\n```\n\nTwo marks on shared axes:\n\n\n```ocaml\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 100 in\n layers [\n line ~x ~y:(Nx.sin x) ~label:\"sin\" ();\n line ~x ~y:(Nx.cos x) ~label:\"cos\" ~line_style:`Dashed ();\n ]\n |> legend |> render_png \"trig.png\"\n```\n\n## Next Steps\n\n- [Getting Started](/docs/hugin/getting-started/) — installation, first plot, key concepts\n- [Marks and Styling](/docs/hugin/marks-and-styling/) — mark catalog, visual properties\n- [Layout and Decorations](/docs/hugin/layout-and-decorations/) — axes, scales, themes, multi-panel\n- [Colors and Colormaps](/docs/hugin/colors-and-colormaps/) — OKLCH colors, palettes, colormaps\n- [Matplotlib Comparison](/docs/hugin/matplotlib-comparison/) — side-by-side with Python\n"
},
{
"path": "packages/hugin/examples/01-line-plot/README.md",
"content": "# Line Plot\n\nCreate data with Nx, build a line plot, and render to PNG in three lines.\n\n\n"
},
{
"path": "packages/hugin/examples/01-line-plot/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets line_plot.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/01-line-plot/main.ml",
"content": "(* Your first plot.\n\n The simplest possible visualization: create data with Nx, build a line plot,\n and render to PNG. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 100 in\n let y = Nx.sin x in\n line ~x ~y () |> render_png \"line_plot.png\"\n"
},
{
"path": "packages/hugin/examples/02-styling/README.md",
"content": "# Styling\n\nMark constructors accept optional visual properties: `~color`, `~line_style`, `~line_width`, `~marker`, and `~alpha`.\n\n\n"
},
{
"path": "packages/hugin/examples/02-styling/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets styling.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/02-styling/main.ml",
"content": "(* Styling.\n\n Every mark constructor accepts optional visual properties as labeled\n arguments. This example shows how to set color, line style, width, and marker\n shape. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 50 in\n let y = Nx.sin x in\n line ~x ~y ~color:Color.vermillion ~line_style:`Dashed ~line_width:2.5\n ~marker:Triangle ~alpha:0.7 ()\n |> render_png \"styling.png\"\n"
},
{
"path": "packages/hugin/examples/03-scatter/README.md",
"content": "# Scatter Plot\n\nThe `point` mark places markers at data coordinates. Pass `~color_by` to map a third variable through the theme's sequential colormap.\n\n\n"
},
{
"path": "packages/hugin/examples/03-scatter/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets scatter.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/03-scatter/main.ml",
"content": "(* Scatter plots.\n\n Point marks place individual markers at data coordinates. Use color_by to map\n a third variable through the theme's sequential colormap. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.rand Nx.float32 [| 200 |] in\n let y = Nx.rand Nx.float32 [| 200 |] in\n let c = Nx.add x y in\n point ~x ~y ~color_by:c ~size:8. ~marker:Circle ()\n |> title \"Random Scatter\" |> render_png \"scatter.png\"\n"
},
{
"path": "packages/hugin/examples/04-bar-chart/README.md",
"content": "# Bar Chart\n\nBar marks draw vertical bars centered at x positions. Use `~xticks` to label the x-axis with category names.\n\n\n"
},
{
"path": "packages/hugin/examples/04-bar-chart/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets bar_chart.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/04-bar-chart/main.ml",
"content": "(* Bar charts.\n\n Bar marks draw vertical bars centered at x positions. Height is measured from\n bottom (default 0). *)\n\nopen Hugin\n\nlet () =\n let x = Nx.create Nx.float32 [| 5 |] [| 1.; 2.; 3.; 4.; 5. |] in\n let h = Nx.create Nx.float32 [| 5 |] [| 4.2; 7.1; 3.8; 9.0; 5.5 |] in\n bar ~x ~height:h ~color:Color.sky_blue ()\n |> title \"Quarterly Revenue\" |> xlabel \"Quarter\" |> ylabel \"Revenue ($M)\"\n |> xticks [ (1., \"Q1\"); (2., \"Q2\"); (3., \"Q3\"); (4., \"Q4\"); (5., \"Q5\") ]\n |> render_png \"bar_chart.png\"\n"
},
{
"path": "packages/hugin/examples/05-histogram/README.md",
"content": "# Histogram\n\nHistogram marks bin continuous data into evenly-spaced intervals. Use `~density:true` to normalize the total area to 1.\n\n\n"
},
{
"path": "packages/hugin/examples/05-histogram/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets histogram.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/05-histogram/main.ml",
"content": "(* Histograms.\n\n Histogram marks bin continuous data. Use ~density:true to normalize so the\n total area equals 1. *)\n\nopen Hugin\n\nlet () =\n let samples = Nx.rand Nx.float32 [| 1000 |] in\n hist ~x:samples ~bins:(`Num 25) ~density:true ~color:Color.green ()\n |> title \"Distribution\" |> xlabel \"Value\" |> render_png \"histogram.png\"\n"
},
{
"path": "packages/hugin/examples/06-layers/README.md",
"content": "# Layers\n\n`layers` overlays multiple marks on shared axes. Any mark with a `~label` automatically appears in the legend.\n\n\n"
},
{
"path": "packages/hugin/examples/06-layers/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets layers.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/06-layers/main.ml",
"content": "(* Layers and legends.\n\n Use layers to overlay different mark types on shared axes. Any mark with a\n ~label automatically appears in the legend. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. 10. 100 in\n let y = Nx.sin x in\n layers\n [\n fill_between ~x ~y1:(Nx.sub_s y 0.3) ~y2:(Nx.add_s y 0.3) ~label:\"± 0.3\"\n ();\n line ~x ~y ~label:\"sin(x)\" ();\n hline ~y:0. ~line_style:`Dashed ~color:Color.gray ~label:\"baseline\" ();\n ]\n |> title \"Sine with Confidence Band\"\n |> xlabel \"x\" |> ylabel \"y\" |> legend |> render_png \"layers.png\"\n"
},
{
"path": "packages/hugin/examples/07-decorations/README.md",
"content": "# Decorations\n\nDecoration functions (`xscale`, `xlim`, `ylim`, `grid_lines`, `xtick_format`) control axis behavior and compose with `|>`.\n\n\n"
},
{
"path": "packages/hugin/examples/07-decorations/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets decorations.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/07-decorations/main.ml",
"content": "(* Axis decorations.\n\n Decorations control axes limits, scales, and grid visibility. They compose\n naturally with the |> pipeline. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 1. 1000. 100 in\n let y = Nx.log x in\n line ~x ~y () |> title \"Logarithmic Scale\" |> xlabel \"x\" |> ylabel \"ln(x)\"\n |> xscale `Log |> xlim 1. 1000. |> ylim 0. 8.\n |> xtick_format (Printf.sprintf \"%.0f\")\n |> grid_lines true\n |> render_png \"decorations.png\"\n"
},
{
"path": "packages/hugin/examples/08-grid-layout/README.md",
"content": "# Grid Layout\n\n`grid` arranges independent plots in rows and columns. Each cell has its own axes and decorations.\n\n\n"
},
{
"path": "packages/hugin/examples/08-grid-layout/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets grid_layout.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/08-grid-layout/main.ml",
"content": "(* Grid layout.\n\n Arrange independent plots in a grid. Each cell is a standalone specification\n with its own axes and decorations. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. (2. *. Float.pi) 100 in\n let p1 = line ~x ~y:(Nx.sin x) () |> title \"sin\" in\n let p2 = line ~x ~y:(Nx.cos x) () |> title \"cos\" in\n let p3 = line ~x ~y:(Nx.tan (Nx.mul_s x 0.3)) () |> title \"tan(0.3x)\" in\n let p4 =\n point ~x ~y:(Nx.sin x) ~color:Color.vermillion ~marker:Plus ()\n |> title \"sin (scatter)\"\n in\n grid [ [ p1; p2 ]; [ p3; p4 ] ] |> render_png \"grid_layout.png\"\n"
},
{
"path": "packages/hugin/examples/09-themes/README.md",
"content": "# Themes\n\nThemes control visual appearance: background, fonts, axes, grid, and data colors. Context functions like `Theme.talk` scale everything up for presentations.\n\n\n"
},
{
"path": "packages/hugin/examples/09-themes/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets themes.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/09-themes/main.ml",
"content": "(* Themes and context scaling.\n\n Themes control the entire visual appearance: colors, fonts, line widths.\n Context functions like Theme.talk scale everything up for presentations. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. 10. 80 in\n let base =\n layers\n [\n line ~x ~y:(Nx.sin x) ~label:\"sin\" ();\n line ~x ~y:(Nx.cos x) ~label:\"cos\" ();\n ]\n |> legend\n in\n grid\n [\n [\n base |> with_theme Theme.default |> title \"Default\";\n base |> with_theme Theme.dark |> title \"Dark\";\n ];\n [\n base |> with_theme Theme.minimal |> title \"Minimal\";\n base |> with_theme (Theme.talk Theme.default) |> title \"Talk\";\n ];\n ]\n |> render_png \"themes.png\"\n"
},
{
"path": "packages/hugin/examples/10-showcase/README.md",
"content": "# Showcase\n\nCombines multiple mark types, layouts, and output formats in a single visualization. Renders to both PNG and SVG.\n\n\n"
},
{
"path": "packages/hugin/examples/10-showcase/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets showcase.png showcase.svg)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/10-showcase/main.ml",
"content": "(* Full showcase.\n\n Demonstrates multiple mark types, layouts, themes, and output formats in a\n single example. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.linspace Nx.float32 0. 10. 100 in\n\n let p1 =\n layers\n [\n line ~x ~y:(Nx.sin x) ~label:\"sin\" ~color:Color.blue ();\n point\n ~x:(Nx.mul_s (Nx.rand Nx.float32 [| 30 |]) 10.)\n ~y:(Nx.sub_s (Nx.mul_s (Nx.rand Nx.float32 [| 30 |]) 2.) 1.)\n ~color:Color.vermillion ~marker:Star ~label:\"noise\" ();\n ]\n |> title \"Lines & Scatter\" |> legend\n in\n\n let p2 =\n let xb = Nx.create Nx.float32 [| 4 |] [| 1.; 2.; 3.; 4. |] in\n let h = Nx.create Nx.float32 [| 4 |] [| 3.; 7.; 2.; 5. |] in\n bar ~x:xb ~height:h ~color:Color.orange () |> title \"Bar Chart\"\n in\n\n let p3 =\n hist ~x:(Nx.rand Nx.float32 [| 500 |]) ~bins:(`Num 20) ~color:Color.green ()\n |> title \"Histogram\"\n in\n\n let p4 =\n let xs = Nx.rand Nx.float32 [| 50 |] in\n let ys = Nx.rand Nx.float32 [| 50 |] in\n let cb = Nx.mul_s xs 100. in\n let sb = Nx.mul_s ys 40. in\n point ~x:xs ~y:ys ~color_by:cb ~size_by:sb ~marker:Circle ()\n |> title \"color_by + size_by\" |> xlabel \"x\" |> ylabel \"y\"\n in\n\n let p5 =\n let xl = Nx.linspace Nx.float32 1. 100. 50 in\n line ~x:xl ~y:(Nx.mul xl xl) ~color:Color.purple ()\n |> title \"Quadratic (log y)\" |> yscale `Log\n in\n\n let p6 =\n let data =\n Nx.init Nx.float32 [| 8; 10 |] (fun idx ->\n let i = Float.of_int idx.(0) and j = Float.of_int idx.(1) in\n Float.sin (i *. 0.5) *. Float.cos (j *. 0.4))\n in\n heatmap ~data ~cmap:Cmap.viridis () |> title \"Heatmap\"\n in\n\n let spec = grid [ [ p1; p2 ]; [ p3; p4 ]; [ p5; p6 ] ] in\n spec |> render_png \"showcase.png\";\n spec |> render_svg \"showcase.svg\"\n"
},
{
"path": "packages/hugin/examples/11-errorbar/README.md",
"content": "# Error Bars\n\n`errorbar` shows measurement uncertainty. Use `` `Symmetric `` for equal +/- errors or `` `Asymmetric `` for independent lower and upper bounds.\n\n\n"
},
{
"path": "packages/hugin/examples/11-errorbar/dune",
"content": "(executable\n (name main)\n (libraries hugin nx))\n\n(rule\n (targets errorbar.png)\n (deps main.exe)\n (action\n (run ./main.exe))\n (mode\n (promote (until-clean))))\n"
},
{
"path": "packages/hugin/examples/11-errorbar/main.ml",
"content": "(* Error bars.\n\n Errorbar marks show measurement uncertainty. Use `Symmetric for equal +/-\n errors or `Asymmetric for independent lower and upper bounds. *)\n\nopen Hugin\n\nlet () =\n let x = Nx.create Nx.float32 [| 6 |] [| 1.; 2.; 3.; 4.; 5.; 6. |] in\n let y = Nx.create Nx.float32 [| 6 |] [| 2.1; 3.8; 3.2; 5.1; 4.5; 6.3 |] in\n let err = Nx.create Nx.float32 [| 6 |] [| 0.3; 0.5; 0.2; 0.6; 0.4; 0.3 |] in\n errorbar ~x ~y ~yerr:(`Symmetric err) ~cap_size:6. ~color:Color.blue ()\n |> title \"Measurements\" |> xlabel \"Trial\" |> ylabel \"Value\"\n |> render_png \"errorbar.png\"\n"
},
{
"path": "packages/hugin/examples/README.md",
"content": "# Hugin Examples\n\nLearn Hugin through progressively complex examples. Start with `01-line-plot`\nand work through the numbered examples in order.\n\n## Examples\n\n| Example | Concept | Key Functions |\n|---------|---------|---------------|\n| [`01-line-plot`](./01-line-plot/) | Your first plot | `line`, `render_png` |\n| [`02-styling`](./02-styling/) | Colors, line styles, markers | `~color`, `~line_style`, `~marker`, `~alpha` |\n| [`03-scatter`](./03-scatter/) | Scatter plots and color mapping | `point`, `~color_by` |\n| [`04-bar-chart`](./04-bar-chart/) | Bar charts with categorical axes | `bar`, `xlabel`, `ylabel`, `xticks` |\n| [`05-histogram`](./05-histogram/) | Histograms and density | `hist`, `~bins`, `~density` |\n| [`06-layers`](./06-layers/) | Overlaying marks and legends | `layers`, `fill_between`, `hline`, `legend` |\n| [`07-decorations`](./07-decorations/) | Axis control and grid lines | `xscale`, `xlim`, `ylim`, `xtick_format`, `grid_lines` |\n| [`08-grid-layout`](./08-grid-layout/) | Multi-panel layouts | `Layout.grid` |\n| [`09-themes`](./09-themes/) | Themes and context scaling | `Theme.default`, `Theme.dark`, `Theme.talk` |\n| [`10-showcase`](./10-showcase/) | Full showcase with multiple outputs | All mark types, `heatmap`, `render_svg` |\n| [`11-errorbar`](./11-errorbar/) | Measurement uncertainty | `errorbar`, `~yerr`, `~cap_size` |\n\n## Running Examples\n\nAll examples can be run with:\n\n```bash\ndune exec dev/hugin/examples//main.exe\n```\n\nFor example:\n\n```bash\ndune exec dev/hugin/examples/01-line-plot/main.exe\n```\n\n## Quick Reference\n\n### Single Plot\n\n```ocaml\nopen Hugin\n\nlet x = Nx.linspace Nx.float32 0. 6.28 100 in\nlet y = Nx.sin x in\nline ~x ~y () |> title \"Sine\" |> render_png \"plot.png\"\n```\n\n### Multiple Marks on Shared Axes\n\n```ocaml\nlayers\n [\n line ~x ~y:(Nx.sin x) ~label:\"sin\" ();\n line ~x ~y:(Nx.cos x) ~label:\"cos\" ~line_style:`Dashed ();\n ]\n|> legend |> render_png \"plot.png\"\n```\n\n### Grid Layout\n\n```ocaml\nlet p1 = line ~x ~y:(Nx.sin x) () |> title \"sin\" in\nlet p2 = line ~x ~y:(Nx.cos x) () |> title \"cos\" in\nLayout.grid [ [ p1; p2 ] ] |> render_png \"grid.png\"\n```\n"
},
{
"path": "packages/hugin/lib/axis.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Per-axis configuration and resolution.\n\n config holds the user-set options (all optional, from decorations). t holds\n the resolved axis with defaults applied and data bounds merged. *)\n\ntype config = {\n label : string option;\n lim : (float * float) option;\n scale : Spec.scale option;\n invert : bool;\n ticks : (float * string) list option;\n tick_format : (float -> string) option;\n}\n\nlet empty_config =\n {\n label = None;\n lim = None;\n scale = None;\n invert = false;\n ticks = None;\n tick_format = None;\n }\n\ntype t = {\n scale : Spec.scale;\n invert : bool;\n lo : float;\n hi : float;\n label : string option;\n ticks : (float * string) list option;\n tick_format : (float -> string) option;\n}\n\nlet resolve ~data_lo ~data_hi (c : config) =\n let scale = Option.value ~default:`Linear c.scale in\n let lo, hi = Option.value ~default:(data_lo, data_hi) c.lim in\n {\n scale;\n invert = c.invert;\n lo;\n hi;\n label = c.label;\n ticks = c.ticks;\n tick_format = c.tick_format;\n }\n\nlet make_scale_and_ticks (a : t) =\n let s = Scale.make ~invert:a.invert a.scale ~lo:a.lo ~hi:a.hi () in\n let ticks =\n match a.ticks with\n | Some t -> t\n | None -> Ticks.generate a.scale ~lo:a.lo ~hi:a.hi ()\n in\n let ticks =\n match a.tick_format with\n | None -> ticks\n | Some f -> List.map (fun (v, _) -> (v, f v)) ticks\n in\n (s, ticks)\n"
},
{
"path": "packages/hugin/lib/axis.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Per-axis configuration and resolution.\n\n {b Internal module.} Consolidates per-axis state into two stages: {!config}\n holds user-set options from decorations (all optional), and {!t} holds the\n resolved axis with defaults applied and data bounds merged. Used by\n {!Prepared} and {!Resolve}. *)\n\n(** {1:config Configuration} *)\n\ntype config = {\n label : string option;\n lim : (float * float) option;\n scale : Spec.scale option;\n invert : bool;\n ticks : (float * string) list option;\n tick_format : (float -> string) option;\n}\n(** The type for per-axis user options collected from decorations. *)\n\nval empty_config : config\n(** [empty_config] is the default configuration: no label, no limits, [invert]\n is [false], scale/ticks/format are [None]. *)\n\n(** {1:resolved Resolved axis} *)\n\ntype t = {\n scale : Spec.scale;\n invert : bool;\n lo : float;\n hi : float;\n label : string option;\n ticks : (float * string) list option;\n tick_format : (float -> string) option;\n}\n(** The type for resolved axes. [scale] defaults to [`Linear], [lo] and [hi]\n come from data bounds unless overridden by {!config.lim}. *)\n\nval resolve : data_lo:float -> data_hi:float -> config -> t\n(** [resolve ~data_lo ~data_hi c] is a resolved axis from [c]. Uses [data_lo]\n and [data_hi] when [c.lim] is [None], and [`Linear] when [c.scale] is\n [None]. *)\n\nval make_scale_and_ticks : t -> Scale.t * (float * string) list\n(** [make_scale_and_ticks a] is [(scale, ticks)] for [a]. Generates ticks via\n {!Ticks.generate} when [a.ticks] is [None], then applies [a.tick_format] if\n set. *)\n"
},
{
"path": "packages/hugin/lib/cairo_backend.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Color helpers *)\n\nlet set_color cr c =\n let r, g, b, a = Color.to_rgba c in\n Ucairo.set_source_rgba cr r g b a\n\nlet set_font cr (font : Theme.font) =\n let weight =\n match font.weight with `Normal -> Ucairo.Normal | `Bold -> Ucairo.Bold\n in\n Ucairo.select_font_face cr font.family weight;\n Ucairo.set_font_size cr font.size\n\n(* Text measurer *)\n\nlet text_measurer cr ~font s =\n set_font cr font;\n let ext = Ucairo.text_extents cr s in\n (ext.width, ext.height)\n\n(* Marker rendering *)\n\nlet draw_marker cr shape size (px, py) =\n let hs = size /. 2. in\n match shape with\n | Spec.Circle ->\n Ucairo.arc cr px py ~r:hs ~a1:0. ~a2:(2. *. Float.pi);\n Ucairo.Path.close cr\n | Spec.Square -> Ucairo.rectangle cr (px -. hs) (py -. hs) ~w:size ~h:size\n | Spec.Triangle ->\n Ucairo.move_to cr px (py -. hs);\n Ucairo.line_to cr (px +. hs) (py +. hs);\n Ucairo.line_to cr (px -. hs) (py +. hs);\n Ucairo.Path.close cr\n | Spec.Plus ->\n Ucairo.move_to cr (px -. hs) py;\n Ucairo.line_to cr (px +. hs) py;\n Ucairo.move_to cr px (py -. hs);\n Ucairo.line_to cr px (py +. hs)\n | Spec.Star ->\n Ucairo.move_to cr (px -. hs) py;\n Ucairo.line_to cr (px +. hs) py;\n Ucairo.move_to cr px (py -. hs);\n Ucairo.line_to cr px (py +. hs);\n let d = hs *. 0.707 in\n Ucairo.move_to cr (px -. d) (py -. d);\n Ucairo.line_to cr (px +. d) (py +. d);\n Ucairo.move_to cr (px +. d) (py -. d);\n Ucairo.line_to cr (px -. d) (py +. d)\n\n(* Primitive rendering *)\n\nlet rec render_primitive cr = function\n | Scene.Path { points; close; fill; stroke; line_width; dash } ->\n if Array.length points < 2 then ()\n else begin\n let x0, y0 = points.(0) in\n Ucairo.move_to cr x0 y0;\n for i = 1 to Array.length points - 1 do\n let x, y = points.(i) in\n Ucairo.line_to cr x y\n done;\n if close then Ucairo.Path.close cr;\n begin match fill with\n | Some c ->\n set_color cr c;\n if stroke <> None then Ucairo.fill_preserve cr else Ucairo.fill cr\n | None -> ()\n end;\n begin match stroke with\n | Some c ->\n set_color cr c;\n Ucairo.set_line_width cr line_width;\n (match dash with\n | [] -> Ucairo.set_dash cr [||]\n | ds -> Ucairo.set_dash cr (Array.of_list ds));\n Ucairo.stroke cr\n | None -> ()\n end\n end\n | Scene.Markers { points; shape; size; sizes; fill; fills; stroke } ->\n let stroke_only =\n match shape with Spec.Plus | Spec.Star -> true | _ -> false\n in\n Array.iteri\n (fun i pt ->\n let s = match sizes with Some ss -> ss.(i) | None -> size in\n let f = match fills with Some fs -> Some fs.(i) | None -> fill in\n Ucairo.Path.clear cr;\n draw_marker cr shape s pt;\n if stroke_only then begin\n let c =\n match f with\n | Some c -> c\n | None -> ( match stroke with Some c -> c | None -> Color.black)\n in\n set_color cr c;\n Ucairo.set_line_width cr (Float.max 1. (s *. 0.15));\n Ucairo.stroke cr\n end\n else begin\n begin match f with\n | Some c ->\n set_color cr c;\n if stroke <> None then Ucairo.fill_preserve cr\n else Ucairo.fill cr\n | None -> ()\n end;\n begin match stroke with\n | Some c ->\n set_color cr c;\n Ucairo.set_line_width cr (Float.max 1. (s *. 0.15));\n Ucairo.stroke cr\n | None -> ()\n end\n end)\n points\n | Scene.Text { x; y; content; font; color; anchor; baseline; angle } ->\n set_font cr font;\n set_color cr color;\n let ext = Ucairo.text_extents cr content in\n let dx =\n match anchor with\n | `Start -> -.ext.x_bearing\n | `Middle -> -.(ext.x_bearing +. (ext.width /. 2.))\n | `End -> -.(ext.x_bearing +. ext.width)\n in\n let dy =\n match baseline with\n | `Top -> -.ext.y_bearing\n | `Middle -> -.(ext.y_bearing +. (ext.height /. 2.))\n | `Bottom -> -.(ext.y_bearing +. ext.height)\n in\n Ucairo.save cr;\n Ucairo.translate cr x y;\n if angle <> 0. then Ucairo.rotate cr angle;\n Ucairo.move_to cr dx dy;\n Ucairo.show_text cr content;\n Ucairo.restore cr\n | Scene.Image { x; y; w; h; data } ->\n let img_surface = Image_util.nx_to_cairo_surface data in\n let img_w = (Nx.shape data).(1) and img_h = (Nx.shape data).(0) in\n Ucairo.save cr;\n Ucairo.translate cr x y;\n Ucairo.scale cr (w /. float img_w) (h /. float img_h);\n Ucairo.set_source_surface cr img_surface ~x:0. ~y:0.;\n Ucairo.paint cr;\n Ucairo.restore cr;\n Ucairo.Surface.finish img_surface\n | Scene.Clip { x; y; w; h; children } ->\n Ucairo.save cr;\n Ucairo.rectangle cr x y ~w ~h;\n Ucairo.clip cr;\n List.iter (render_primitive cr) children;\n Ucairo.restore cr\n | Scene.Group children -> List.iter (render_primitive cr) children\n\n(* Scene rendering *)\n\nlet render_scene cr (scene : Scene.t) =\n Ucairo.set_antialias cr Ucairo.Antialias_default;\n Ucairo.set_line_cap cr Ucairo.Round;\n Ucairo.set_line_join cr Ucairo.Join_round;\n List.iter (render_primitive cr) scene.primitives\n\n(* Entry points *)\n\nlet render_to_png filename ~width ~height (scene : Scene.t) =\n let w = int_of_float width and h = int_of_float height in\n let surface = Ucairo.Image.create ~w ~h in\n let cr = Ucairo.create surface in\n render_scene cr scene;\n Ucairo.Png.write surface filename;\n Ucairo.Surface.finish surface\n\nlet render_to_pdf filename ~width ~height (scene : Scene.t) =\n let surface = Ucairo.Pdf.create filename ~w:width ~h:height in\n let cr = Ucairo.create surface in\n render_scene cr scene;\n Ucairo.Surface.finish surface\n\nlet render_to_buffer ~width ~height (scene : Scene.t) =\n let w = int_of_float width and h = int_of_float height in\n let surface = Ucairo.Image.create ~w ~h in\n let cr = Ucairo.create surface in\n render_scene cr scene;\n let buf = Buffer.create 4096 in\n Ucairo.Png.write_to_stream surface (Buffer.add_string buf);\n Ucairo.Surface.finish surface;\n Buffer.contents buf\n\nlet show_interactive ~theme ~width ~height prepared =\n let w = int_of_float width and h = int_of_float height in\n let csdl = Cairo_sdl.create ~width:w ~height:h ~title:\"Hugin\" in\n\n let render_current () =\n let cr = Cairo_sdl.context csdl in\n let cw = float (Cairo_sdl.width csdl) in\n let ch = float (Cairo_sdl.height csdl) in\n let tm = text_measurer cr in\n let scene =\n Resolve.resolve_prepared ~text_measurer:tm ~theme ~width:cw ~height:ch\n prepared\n in\n render_scene cr scene;\n Cairo_sdl.present csdl\n in\n\n render_current ();\n\n let ev = Usdl.Event.create () in\n let quit = ref false in\n while not !quit do\n if not (Usdl.Event.wait ev) then quit := true\n else\n begin match Usdl.Event.typ ev with\n | `Quit -> quit := true\n | `Window_event ->\n begin match Usdl.Event.window_event_id ev with\n | `Resized | `Size_changed ->\n Cairo_sdl.resize csdl;\n render_current ()\n | `Exposed -> render_current ()\n | `Close -> quit := true\n | _ -> ()\n end\n | `Key_down ->\n let keycode = Usdl.Event.keycode ev in\n if keycode = Usdl.Keycode.escape || keycode = Usdl.Keycode.q then\n quit := true\n | _ -> ()\n end\n done;\n Cairo_sdl.destroy csdl\n"
},
{
"path": "packages/hugin/lib/cairo_backend.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Cairo rendering backend.\n\n {b Internal module.} Renders {!Scene.t} to PNG, PDF, or an interactive SDL\n window via Cairo. *)\n\n(** {1:measurer Text measurement} *)\n\nval text_measurer : Ucairo.t -> Resolve.text_measurer\n(** [text_measurer cr] is a text measurer backed by {!Ucairo.text_extents}. *)\n\n(** {1:rendering Rendering} *)\n\nval render_scene : Ucairo.t -> Scene.t -> unit\n(** [render_scene cr scene] draws [scene] onto [cr]. *)\n\nval render_to_png : string -> width:float -> height:float -> Scene.t -> unit\n(** [render_to_png filename ~width ~height scene] writes [scene] as a PNG image.\n*)\n\nval render_to_pdf : string -> width:float -> height:float -> Scene.t -> unit\n(** [render_to_pdf filename ~width ~height scene] writes [scene] as a\n single-page PDF. *)\n\nval render_to_buffer : width:float -> height:float -> Scene.t -> string\n(** [render_to_buffer ~width ~height scene] is the PNG-encoded contents of\n [scene] as a string. *)\n\n(** {1:interactive Interactive display} *)\n\nval show_interactive :\n theme:Theme.t -> width:float -> height:float -> Prepared.t -> unit\n(** [show_interactive ~theme ~width ~height prepared] opens an SDL window and\n renders [prepared]. Compiles data once; only re-resolves layout on resize.\n Exits on Escape, Q, or window close. *)\n"
},
{
"path": "packages/hugin/lib/cairo_sdl.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Cairo-SDL integration: shared ARGB8888 surface *)\n\ntype t = {\n window : Usdl.Window.t;\n renderer : Usdl.Renderer.t;\n mutable surface : Usdl.Surface.t;\n mutable cairo_surface : Ucairo.surface;\n mutable context : Ucairo.t;\n mutable width : int;\n mutable height : int;\n}\n\nlet make_cairo_context surface =\n let pixels = Usdl.Surface.pixels surface in\n let stride = Usdl.Surface.pitch surface in\n let total = Bigarray.Array1.dim pixels in\n let h = total / stride in\n let w = stride / 4 in\n let cs = Ucairo.Image.create_for_data8 pixels ~w ~h ~stride in\n let cr = Ucairo.create cs in\n (cs, cr, w, h)\n\nlet create ~width ~height ~title =\n Usdl.init ();\n let window = Usdl.Window.create ~title ~w:width ~h:height in\n let renderer = Usdl.Renderer.create window in\n let ow, oh = Usdl.Renderer.output_size renderer in\n let surface = Usdl.Surface.create_argb8888 ~w:ow ~h:oh in\n let cairo_surface, context, w, h = make_cairo_context surface in\n { window; renderer; surface; cairo_surface; context; width = w; height = h }\n\nlet context t = t.context\nlet width t = t.width\nlet height t = t.height\n\nlet present t =\n Ucairo.Surface.flush t.cairo_surface;\n let tex = Usdl.Texture.of_surface t.renderer t.surface in\n Usdl.Renderer.clear t.renderer;\n Usdl.Renderer.copy t.renderer tex;\n Usdl.Renderer.present t.renderer;\n Usdl.Texture.destroy tex\n\nlet resize t =\n let nw, nh = Usdl.Renderer.output_size t.renderer in\n if nw <> t.width || nh <> t.height then\n begin if nw > 0 && nh > 0 then begin\n Ucairo.Surface.finish t.cairo_surface;\n Usdl.Surface.destroy t.surface;\n let surface = Usdl.Surface.create_argb8888 ~w:nw ~h:nh in\n let cairo_surface, context, w, h = make_cairo_context surface in\n t.surface <- surface;\n t.cairo_surface <- cairo_surface;\n t.context <- context;\n t.width <- w;\n t.height <- h\n end\n end\n\nlet destroy t =\n Ucairo.Surface.finish t.cairo_surface;\n Usdl.Surface.destroy t.surface;\n Usdl.Renderer.destroy t.renderer;\n Usdl.Window.destroy t.window;\n Usdl.quit ()\n"
},
{
"path": "packages/hugin/lib/cairo_sdl.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Cairo-SDL integration.\n\n {b Internal module.} Manages a shared ARGB8888 surface between Cairo and SDL\n for interactive rendering. *)\n\ntype t\n(** The type for Cairo-SDL contexts. *)\n\nval create : width:int -> height:int -> title:string -> t\n(** [create ~width ~height ~title] initializes SDL, creates a resizable window,\n and sets up a shared Cairo surface.\n\n Raises [Failure] if SDL initialization fails. *)\n\nval context : t -> Ucairo.t\n(** [context t] is the current Cairo drawing context. Valid until the next\n {!present} or {!resize}. *)\n\nval width : t -> int\n(** [width t] is the current surface width in pixels. *)\n\nval height : t -> int\n(** [height t] is the current surface height in pixels. *)\n\nval present : t -> unit\n(** [present t] flushes the Cairo surface to the SDL window and prepares a fresh\n Cairo context for the next frame. *)\n\nval resize : t -> unit\n(** [resize t] updates the surface dimensions to match the renderer output size.\n No-op if the size has not changed. *)\n\nval destroy : t -> unit\n(** [destroy t] frees all SDL and Cairo resources and calls {!Usdl.quit}. *)\n"
},
{
"path": "packages/hugin/lib/cmap.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype t = Color.t array\n\nlet eval t v =\n let v = Float.max 0. (Float.min 1. v) in\n let i = int_of_float (v *. 255.) in\n t.(min i 255)\n\nlet of_colors stops =\n let n = Array.length stops in\n if n < 2 then invalid_arg \"Cmap.of_colors: need at least 2 stops\";\n Array.init 256 (fun i ->\n let v = float i /. 255. in\n let scaled = v *. float (n - 1) in\n let idx = int_of_float scaled in\n let idx = min idx (n - 2) in\n let frac = scaled -. float idx in\n Color.mix frac stops.(idx) stops.(idx + 1))\n\n(* Decode a canonical 256-entry hex-encoded colormap *)\n\nlet hex_digit c =\n match c with\n | '0' .. '9' -> Char.code c - Char.code '0'\n | 'a' .. 'f' -> 10 + Char.code c - Char.code 'a'\n | 'A' .. 'F' -> 10 + Char.code c - Char.code 'A'\n | _ -> invalid_arg (Printf.sprintf \"Cmap.hex_digit: invalid hex digit %C\" c)\n\nlet decode_hex_cmap hex =\n Array.init 256 (fun i ->\n let off = i * 6 in\n let byte j =\n let h = hex_digit (String.unsafe_get hex (off + (j * 2))) in\n let l = hex_digit (String.unsafe_get hex (off + (j * 2) + 1)) in\n float ((h lsl 4) lor l) /. 255.\n in\n Color.rgb ~r:(byte 0) ~g:(byte 1) ~b:(byte 2) ())\n\nlet viridis = decode_hex_cmap Cmap_data.viridis_hex\nlet plasma = decode_hex_cmap Cmap_data.plasma_hex\nlet inferno = decode_hex_cmap Cmap_data.inferno_hex\nlet magma = decode_hex_cmap Cmap_data.magma_hex\nlet cividis = decode_hex_cmap Cmap_data.cividis_hex\nlet coolwarm = decode_hex_cmap Cmap_data.coolwarm_hex\n\nlet gray =\n of_colors [| Color.rgb ~r:0. ~g:0. ~b:0. (); Color.rgb ~r:1. ~g:1. ~b:1. () |]\n\nlet gray_r =\n of_colors [| Color.rgb ~r:1. ~g:1. ~b:1. (); Color.rgb ~r:0. ~g:0. ~b:0. () |]\n\nlet hot =\n of_colors\n [|\n Color.rgb ~r:0. ~g:0. ~b:0. ();\n Color.rgb ~r:0.7 ~g:0. ~b:0. ();\n Color.rgb ~r:1. ~g:0.6 ~b:0. ();\n Color.rgb ~r:1. ~g:1. ~b:1. ();\n |]\n"
},
{
"path": "packages/hugin/lib/cmap.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Colormaps.\n\n A colormap is a continuous mapping from \\[[0];[1]\\] to {!Color.t}.\n Internally stored as a 256-entry lookup table with OKLCH interpolation, so\n {!eval} is a single array access. *)\n\n(** {1:types Types} *)\n\ntype t\n(** The type for colormaps. *)\n\n(** {1:eval Evaluation} *)\n\nval eval : t -> float -> Color.t\n(** [eval cmap v] is the color at position [v], clamped to \\[[0];[1]\\]. *)\n\n(** {1:constructors Constructors} *)\n\nval of_colors : Color.t array -> t\n(** [of_colors stops] is a colormap interpolating linearly through [stops] in\n OKLCH space. The stops are evenly spaced from [0] to [1].\n\n Raises [Invalid_argument] if [stops] has fewer than 2 elements. *)\n\n(** {1:predefined Predefined colormaps}\n\n Perceptually uniform sequential colormaps from the\n {{:https://bids.github.io/colormap/}viridis family}, plus a diverging\n colormap. *)\n\nval viridis : t\nval plasma : t\nval inferno : t\nval magma : t\nval cividis : t\nval coolwarm : t\n\nval gray : t\n(** Linear grayscale (black to white). *)\n\nval gray_r : t\n(** Reversed grayscale (white to black). The standard default for astronomical\n image display. *)\n\nval hot : t\n(** Black-red-yellow-white. Common in X-ray astronomy. *)\n"
},
{
"path": "packages/hugin/lib/cmap_data.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Canonical 256-entry colormap data encoded as hex strings. Each string is 1536\n characters: 256 entries of 6 hex chars (RRGGBB). *)\n\nlet viridis_hex =\n \"44015444025645045745055946075a46085c460a5d460b5e470d60470e6147106347116447136548146748166848176948186a481a6c481b6d481c6e481d6f481f70482071482173482374482475482576482677482878482979472a7a472c7a472d7b472e7c472f7d46307e46327e46337f463480453581453781453882443983443a83443b84433d84433e85423f854240864241864142874144874045884046883f47883f48893e49893e4a893e4c8a3d4d8a3d4e8a3c4f8a3c508b3b518b3b528b3a538b3a548c39558c39568c38588c38598c375a8c375b8d365c8d365d8d355e8d355f8d34608d34618d33628d33638d32648e32658e31668e31678e31688e30698e306a8e2f6b8e2f6c8e2e6d8e2e6e8e2e6f8e2d708e2d718e2c718e2c728e2c738e2b748e2b758e2a768e2a778e2a788e29798e297a8e297b8e287c8e287d8e277e8e277f8e27808e26818e26828e26828e25838e25848e25858e24868e24878e23888e23898e238a8d228b8d228c8d228d8d218e8d218f8d21908d21918c20928c20928c20938c1f948c1f958b1f968b1f978b1f988b1f998a1f9a8a1e9b8a1e9c891e9d891f9e891f9f881fa0881fa1881fa1871fa28720a38620a48621a58521a68522a78522a88423a98324aa8325ab8225ac8226ad8127ad8128ae8029af7f2ab07f2cb17e2db27d2eb37c2fb47c31b57b32b67a34b67935b77937b87838b9773aba763bbb753dbc743fbc7340bd7242be7144bf7046c06f48c16e4ac16d4cc26c4ec36b50c46a52c56954c56856c66758c7655ac8645cc8635ec96260ca6063cb5f65cb5e67cc5c69cd5b6ccd5a6ece5870cf5773d05675d05477d1537ad1517cd2507fd34e81d34d84d44b86d54989d5488bd6468ed64590d74393d74195d84098d83e9bd93c9dd93ba0da39a2da37a5db36a8db34aadc32addc30b0dd2fb2dd2db5de2bb8de29bade28bddf26c0df25c2df23c5e021c8e020cae11fcde11dd0e11cd2e21bd5e21ad8e219dae319dde318dfe318e2e418e5e419e7e419eae51aece51befe51cf1e51df4e61ef6e620f8e621fbe723fde725\"\n\nlet plasma_hex =\n \"0d088710078813078916078a19068c1b068d1d068e20068f2206902406912605912805922a05932c05942e05952f059631059733059735049837049938049a3a049a3c049b3e049c3f049c41049d43039e44039e46039f48039f4903a04b03a14c02a14e02a25002a25102a35302a35502a45601a45801a45901a55b01a55c01a65e01a66001a66100a76300a76400a76600a76700a86900a86a00a86c00a86e00a86f00a87100a87201a87401a87501a87701a87801a87a02a87b02a87d03a87e03a88004a88104a78305a78405a78606a68707a68808a68a09a58b0aa58d0ba58e0ca48f0da4910ea3920fa39410a29511a19613a19814a099159f9a169f9c179e9d189d9e199da01a9ca11b9ba21d9aa31e9aa51f99a62098a72197a82296aa2395ab2494ac2694ad2793ae2892b02991b12a90b22b8fb32c8eb42e8db52f8cb6308bb7318ab83289ba3388bb3488bc3587bd3786be3885bf3984c03a83c13b82c23c81c33d80c43e7fc5407ec6417dc7427cc8437bc9447aca457acb4679cc4778cc4977cd4a76ce4b75cf4c74d04d73d14e72d24f71d35171d45270d5536fd5546ed6556dd7566cd8576bd9586ada5a6ada5b69db5c68dc5d67dd5e66de5f65de6164df6263e06363e16462e26561e26660e3685fe4695ee56a5de56b5de66c5ce76e5be76f5ae87059e97158e97257ea7457eb7556eb7655ec7754ed7953ed7a52ee7b51ef7c51ef7e50f07f4ff0804ef1814df1834cf2844bf3854bf3874af48849f48948f58b47f58c46f68d45f68f44f79044f79143f79342f89441f89540f9973ff9983ef99a3efa9b3dfa9c3cfa9e3bfb9f3afba139fba238fca338fca537fca636fca835fca934fdab33fdac33fdae32fdaf31fdb130fdb22ffdb42ffdb52efeb72dfeb82cfeba2cfebb2bfebd2afebe2afec029fdc229fdc328fdc527fdc627fdc827fdca26fdcb26fccd25fcce25fcd025fcd225fbd324fbd524fbd724fad824fada24f9dc24f9dd25f8df25f8e125f7e225f7e425f6e626f6e826f5e926f5eb27f4ed27f3ee27f3f027f2f227f1f426f1f525f0f724f0f921\"\n\nlet inferno_hex =\n \"00000401000501010601010802010a02020c02020e03021004031204031405041706041907051b08051d09061f0a07220b07240c08260d08290e092b10092d110a30120a32140b34150b37160b39180c3c190c3e1b0c411c0c431e0c451f0c48210c4a230c4c240c4f260c51280b53290b552b0b572d0b592f0a5b310a5c320a5e340a5f3609613809623909633b09643d09653e0966400a67420a68440a68450a69470b6a490b6a4a0c6b4c0c6b4d0d6c4f0d6c510e6c520e6d540f6d550f6d57106e59106e5a116e5c126e5d126e5f136e61136e62146e64156e65156e67166e69166e6a176e6c186e6d186e6f196e71196e721a6e741a6e751b6e771c6d781c6d7a1d6d7c1d6d7d1e6d7f1e6c801f6c82206c84206b85216b87216b88226a8a226a8c23698d23698f24699025689225689326679526679727669827669a28659b29649d29649f2a63a02a63a22b62a32c61a52c60a62d60a82e5fa92e5eab2f5ead305dae305cb0315bb1325ab3325ab43359b63458b73557b93556ba3655bc3754bd3853bf3952c03a51c13a50c33b4fc43c4ec63d4dc73e4cc83f4bca404acb4149cc4248ce4347cf4446d04545d24644d34743d44842d54a41d74b3fd84c3ed94d3dda4e3cdb503bdd513ade5238df5337e05536e15635e25734e35933e45a31e55c30e65d2fe75e2ee8602de9612bea632aeb6429eb6628ec6726ed6925ee6a24ef6c23ef6e21f06f20f1711ff1731df2741cf3761bf37819f47918f57b17f57d15f67e14f68013f78212f78410f8850ff8870ef8890cf98b0bf98c0af98e09fa9008fa9207fa9407fb9606fb9706fb9906fb9b06fb9d07fc9f07fca108fca309fca50afca60cfca80dfcaa0ffcac11fcae12fcb014fcb216fcb418fbb61afbb81dfbba1ffbbc21fbbe23fac026fac228fac42afac62df9c72ff9c932f9cb35f8cd37f8cf3af7d13df7d340f6d543f6d746f5d949f5db4cf4dd4ff4df53f4e156f3e35af3e55df2e661f2e865f2ea69f1ec6df1ed71f1ef75f1f179f2f27df2f482f3f586f3f68af4f88ef5f992f6fa96f8fb9af9fc9dfafda1fcffa4\"\n\nlet magma_hex =\n \"00000401000501010601010802010902020b02020d03030f03031204041405041606051806051a07061c08071e0907200a08220b09240c09260d0a290e0b2b100b2d110c2f120d31130d34140e36150e38160f3b180f3d19103f1a10421c10441d11471e114920114b21114e22115024125325125527125829115a2a115c2c115f2d11612f116331116533106734106936106b38106c390f6e3b0f703d0f713f0f72400f74420f75440f764510774710784910784a10794c117a4e117b4f127b51127c52137c54137d56147d57157e59157e5a167e5c167f5d177f5f187f601880621980641a80651a80671b80681c816a1c816b1d816d1d816e1e81701f81721f817320817521817621817822817922827b23827c23827e24828025828125818326818426818627818827818928818b29818c29818e2a81902a81912b81932b80942c80962c80982d80992d809b2e7f9c2e7f9e2f7fa02f7fa1307ea3307ea5317ea6317da8327daa337dab337cad347cae347bb0357bb2357bb3367ab5367ab73779b83779ba3878bc3978bd3977bf3a77c03a76c23b75c43c75c53c74c73d73c83e73ca3e72cc3f71cd4071cf4070d0416fd2426fd3436ed5446dd6456cd8456cd9466bdb476adc4869de4968df4a68e04c67e24d66e34e65e44f64e55064e75263e85362e95462ea5661eb5760ec5860ed5a5fee5b5eef5d5ef05f5ef1605df2625df2645cf3655cf4675cf4695cf56b5cf66c5cf66e5cf7705cf7725cf8745cf8765cf9785df9795df97b5dfa7d5efa7f5efa815ffb835ffb8560fb8761fc8961fc8a62fc8c63fc8e64fc9065fd9266fd9467fd9668fd9869fd9a6afd9b6bfe9d6cfe9f6dfea16efea36ffea571fea772fea973feaa74feac76feae77feb078feb27afeb47bfeb67cfeb77efeb97ffebb81febd82febf84fec185fec287fec488fec68afec88cfeca8dfecc8ffecd90fecf92fed194fed395fed597fed799fed89afdda9cfddc9efddea0fde0a1fde2a3fde3a5fde5a7fde7a9fde9aafdebacfcecaefceeb0fcf0b2fcf2b4fcf4b6fcf6b8fcf7b9fcf9bbfcfbbdfcfdbf\"\n\nlet cividis_hex =\n \"00224e00234f00245100255300255400265600275800285900285b00295d002a5f002a61002b62002c64002c66002d68002e6a002e6c002f6d00306f0030700031700031710132710533710833700c34700f357012357014367016377018376f1a386f1c396f1e3a6f203a6f213b6e233c6e243c6e263d6e273e6e293f6e2a3f6d2b406d2d416d2e416d2f426d31436d32436d33446d34456c35456c36466c38476c39486c3a486c3b496c3c4a6c3d4a6c3e4b6c3f4c6c404c6c414d6c424e6c434e6c444f6c45506c46516c47516c48526c49536c4a536c4b546c4c556c4d556c4e566c4f576c50576c51586d52596d535a6d545a6d555b6d555c6d565c6d575d6d585e6d595e6e5a5f6e5b606e5c616e5d616e5e626e5e636f5f636f60646f61656f62656f636670646770656870656870666970676a71686a71696b716a6c716b6d726c6d726c6e726d6f726e6f736f70737071737172747272747273747374757474757575757676767777767777777878777979777a7a787b7a787c7b787d7c787e7c787e7d787f7e78807f78817f788280798381798482798582798683798784788885788985788a86788b87788c88788d88788e89788f8a78908b78918b78928c78928d78938e78948e77958f779690779791779892779992779a93769b94769c95769d95769e96769f9775a09875a19975a29975a39a74a49b74a59c74a69c74a79d73a89e73a99f73aaa073aba072aca172ada272aea371afa471b0a571b1a570b3a670b4a76fb5a86fb6a96fb7a96eb8aa6eb9ab6dbaac6dbbad6dbcae6cbdae6cbeaf6bbfb06bc0b16ac1b26ac2b369c3b369c4b468c5b568c6b667c7b767c8b866c9b965cbb965ccba64cdbb63cebc63cfbd62d0be62d1bf61d2c060d3c05fd4c15fd5c25ed6c35dd7c45cd9c55cdac65bdbc75adcc859ddc858dec958dfca57e0cb56e1cc55e2cd54e4ce53e5cf52e6d051e7d150e8d24fe9d34eead34cebd44bedd54aeed649efd748f0d846f1d945f2da44f3db42f5dc41f6dd3ff7de3ef8df3cf9e03afbe138fce236fde334fee434fee535fee636fee838\"\n\nlet coolwarm_hex =\n \"3b4cc03c4ec23d50c33e51c53f53c64055c84257c94358cb445acc455cce465ecf485fd14961d24a63d34b64d54c66d64e68d84f69d9506bda516ddb536edd5470de5572df5673e05875e15977e35a78e45b7ae55d7ce65e7de75f7fe86180e96282ea6384eb6485ec6687ed6788ee688aef6a8bef6b8df06c8ff16e90f26f92f37093f37295f47396f57597f67699f6779af7799cf87a9df87b9ff97da0f97ea1fa80a3fa81a4fb82a6fb84a7fc85a8fc86a9fc88abfd89acfd8badfd8caffe8db0fe8fb1fe90b2fe92b4fe93b5fe94b6ff96b7ff97b8ff98b9ff9abbff9bbcff9dbdff9ebeff9fbfffa1c0ffa2c1ffa3c2fea5c3fea6c4fea7c5fea9c6fdaac7fdabc8fdadc9fdaec9fcafcafcb1cbfcb2ccfbb3cdfbb5cdfab6cefab7cff9b9d0f9bad0f8bbd1f8bcd2f7bed2f6bfd3f6c0d4f5c1d4f4c3d5f4c4d5f3c5d6f2c6d6f1c7d7f0c9d7f0cad8efcbd8eeccd9edcdd9eccedaebcfdaead1dae9d2dbe8d3dbe7d4dbe6d5dbe5d6dce4d7dce3d8dce2d9dce1dadce0dbdcdedcdddddddcdcdedcdbdfdbd9e0dbd8e1dad6e2dad5e3d9d3e4d9d2e5d8d1e6d7cfe7d7cee8d6cce9d5cbead5c9ead4c8ebd3c6ecd3c5edd2c3edd1c2eed0c0efcfbfefcebdf0cdbbf1cdbaf1ccb8f2cbb7f2cab5f2c9b4f3c8b2f3c7b1f4c6aff4c5adf5c4acf5c2aaf5c1a9f5c0a7f6bfa6f6bea4f6bda2f7bca1f7ba9ff7b99ef7b89cf7b79bf7b599f7b497f7b396f7b194f7b093f7af91f7ad90f7ac8ef7aa8cf7a98bf7a889f7a688f6a586f6a385f6a283f5a081f59f80f59d7ef59c7df49a7bf4987af39778f39577f39475f29274f29072f18f71f18d6ff08b6ef08a6cef886bee8669ee8468ed8366ec8165ec7f63eb7d62ea7b60e97a5fe9785de8765ce7745be67259e57058e46e56e36c55e36b54e26952e16751e0654fdf634ede614ddd5f4bdc5d4ada5a49d95847d85646d75445d65244d55042d44e41d24b40d1493fd0473dcf453ccd423bcc403acb3e38ca3b37c83836c73635c53334c43032c32e31c12b30c0282fbe242ebd1f2dbb1b2cba162bb8122ab70d28b50927b40426\"\n"
},
{
"path": "packages/hugin/lib/color.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype t = { l : float; c : float; h : float; a : float }\n\n(* Constructors *)\n\nlet oklch ~l ~c ~h () = { l; c; h; a = 1. }\nlet oklcha ~l ~c ~h ~a () = { l; c; h; a }\n\n(* sRGB <-> linear RGB *)\n\nlet srgb_to_linear c =\n if c <= 0.04045 then c /. 12.92 else Float.pow ((c +. 0.055) /. 1.055) 2.4\n\nlet linear_to_srgb c =\n if c <= 0.0031308 then 12.92 *. c\n else (1.055 *. Float.pow c (1. /. 2.4)) -. 0.055\n\n(* Linear RGB -> OKLab *)\n\nlet linear_rgb_to_oklab r g b =\n let l = (0.4122214708 *. r) +. (0.5363325363 *. g) +. (0.0514459929 *. b) in\n let m = (0.2119034982 *. r) +. (0.6806995451 *. g) +. (0.1073969566 *. b) in\n let s = (0.0883024619 *. r) +. (0.2164557896 *. g) +. (0.6898418685 *. b) in\n let l = Float.cbrt l and m = Float.cbrt m and s = Float.cbrt s in\n let lab_l =\n (0.2104542553 *. l) +. (0.7936177850 *. m) -. (0.0040720468 *. s)\n in\n let lab_a =\n (1.9779984951 *. l) -. (2.4285922050 *. m) +. (0.4505937099 *. s)\n in\n let lab_b =\n (0.0259040371 *. l) +. (0.7827717662 *. m) -. (0.8086757660 *. s)\n in\n (lab_l, lab_a, lab_b)\n\n(* OKLab -> linear RGB *)\n\nlet oklab_to_linear_rgb lab_l lab_a lab_b =\n let l = lab_l +. (0.3963377774 *. lab_a) +. (0.2158037573 *. lab_b) in\n let m = lab_l -. (0.1055613458 *. lab_a) -. (0.0638541728 *. lab_b) in\n let s = lab_l -. (0.0894841775 *. lab_a) -. (1.2914855480 *. lab_b) in\n let l = l *. l *. l and m = m *. m *. m and s = s *. s *. s in\n let r = (4.0767416621 *. l) -. (3.3077115913 *. m) +. (0.2309699292 *. s) in\n let g = (-1.2684380046 *. l) +. (2.6097574011 *. m) -. (0.3413193965 *. s) in\n let b = (-0.0041960863 *. l) -. (0.7034186147 *. m) +. (1.7076147010 *. s) in\n (r, g, b)\n\n(* OKLab <-> OKLCH *)\n\nlet oklab_to_oklch lab_l lab_a lab_b =\n let c = Float.sqrt ((lab_a *. lab_a) +. (lab_b *. lab_b)) in\n let h = Float.atan2 lab_b lab_a *. 180. /. Float.pi in\n let h = if h < 0. then h +. 360. else h in\n (lab_l, c, h)\n\nlet oklch_to_oklab l c h =\n let h_rad = h *. Float.pi /. 180. in\n (l, c *. Float.cos h_rad, c *. Float.sin h_rad)\n\n(* sRGB -> OKLCH *)\n\nlet of_srgb r g b =\n let lr = srgb_to_linear r\n and lg = srgb_to_linear g\n and lb = srgb_to_linear b in\n let lab_l, lab_a, lab_b = linear_rgb_to_oklab lr lg lb in\n oklab_to_oklch lab_l lab_a lab_b\n\n(* OKLCH -> sRGB *)\n\nlet to_srgb l c h =\n let lab_l, lab_a, lab_b = oklch_to_oklab l c h in\n let lr, lg, lb = oklab_to_linear_rgb lab_l lab_a lab_b in\n let clamp v = Float.max 0. (Float.min 1. v) in\n ( linear_to_srgb (clamp lr),\n linear_to_srgb (clamp lg),\n linear_to_srgb (clamp lb) )\n\nlet rgb ~r ~g ~b () =\n let l, c, h = of_srgb r g b in\n { l; c; h; a = 1. }\n\nlet rgba ~r ~g ~b ~a () =\n let l, c, h = of_srgb r g b in\n { l; c; h; a }\n\nlet hex_digit c =\n match c with\n | '0' .. '9' -> Char.code c - Char.code '0'\n | 'a' .. 'f' -> Char.code c - Char.code 'a' + 10\n | 'A' .. 'F' -> Char.code c - Char.code 'A' + 10\n | _ -> invalid_arg (Printf.sprintf \"Color.hex: invalid hex digit %C\" c)\n\nlet hex_byte s i =\n let hi = hex_digit (String.get s i) in\n let lo = hex_digit (String.get s (i + 1)) in\n float ((hi * 16) + lo) /. 255.\n\nlet hex s =\n let n = String.length s in\n let off = if n > 0 && String.get s 0 = '#' then 1 else 0 in\n let len = n - off in\n match len with\n | 6 ->\n let r = hex_byte s off\n and g = hex_byte s (off + 2)\n and b = hex_byte s (off + 4) in\n rgb ~r ~g ~b ()\n | 8 ->\n let r = hex_byte s off and g = hex_byte s (off + 2) in\n let b = hex_byte s (off + 4) and a = hex_byte s (off + 6) in\n rgba ~r ~g ~b ~a ()\n | _ ->\n invalid_arg\n (Printf.sprintf \"Color.hex: expected 6 or 8 hex digits, got %d\" len)\n\n(* Accessors *)\n\nlet lightness t = t.l\nlet chroma t = t.c\nlet hue t = t.h\nlet alpha t = t.a\n\n(* Converting *)\n\nlet to_rgba t =\n let r, g, b = to_srgb t.l t.c t.h in\n (r, g, b, t.a)\n\n(* Operations *)\n\nlet with_alpha a t = { t with a }\nlet lighten amount t = { t with l = Float.min 1. (t.l +. amount) }\nlet darken amount t = { t with l = Float.max 0. (t.l -. amount) }\n\nlet interpolate_hue ratio h1 h2 =\n let diff = h2 -. h1 in\n let diff =\n if diff > 180. then diff -. 360.\n else if diff < -180. then diff +. 360.\n else diff\n in\n let h = h1 +. (ratio *. diff) in\n if h < 0. then h +. 360. else if h >= 360. then h -. 360. else h\n\nlet mix ratio a b =\n {\n l = a.l +. (ratio *. (b.l -. a.l));\n c = a.c +. (ratio *. (b.c -. a.c));\n h = interpolate_hue ratio a.h b.h;\n a = a.a +. (ratio *. (b.a -. a.a));\n }\n\n(* Named colors — Okabe-Ito *)\n\nlet orange = hex \"#E69F00\"\nlet sky_blue = hex \"#56B4E9\"\nlet green = hex \"#009E73\"\nlet yellow = hex \"#F0E442\"\nlet blue = hex \"#0072B2\"\nlet vermillion = hex \"#D55E00\"\nlet purple = hex \"#CC79A7\"\nlet black = { l = 0.; c = 0.; h = 0.; a = 1. }\nlet white = { l = 1.; c = 0.; h = 0.; a = 1. }\nlet gray = oklch ~l:0.5 ~c:0. ~h:0. ()\n\n(* Formatting *)\n\nlet pp fmt t = Format.fprintf fmt \"oklch(%g %g %g / %g)\" t.l t.c t.h t.a\n"
},
{
"path": "packages/hugin/lib/color.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Perceptually uniform colors.\n\n Colors are represented internally in the\n {{:https://bottosson.github.io/posts/oklab/}OKLCH} color space. All\n operations ({!lighten}, {!darken}, {!mix}) produce perceptually uniform\n results: equal numerical steps yield equal perceived differences.\n\n Constructors accept common input formats (sRGB, hex) and convert to OKLCH on\n creation. The reverse conversion {!to_rgba} is called only at render time.\n*)\n\n(** {1:types Types} *)\n\ntype t\n(** The type for colors in OKLCH space. Components are lightness \\[0, 1\\],\n chroma \\[0, ~0.4\\], hue \\[0, 360), and alpha \\[0, 1\\]. *)\n\n(** {1:constructors Constructors} *)\n\nval oklch : l:float -> c:float -> h:float -> unit -> t\n(** [oklch ~l ~c ~h ()] is the fully opaque OKLCH color with lightness [l],\n chroma [c], and hue [h] (in degrees). *)\n\nval oklcha : l:float -> c:float -> h:float -> a:float -> unit -> t\n(** [oklcha ~l ~c ~h ~a ()] is like {!oklch} with alpha [a]. *)\n\nval rgb : r:float -> g:float -> b:float -> unit -> t\n(** [rgb ~r ~g ~b ()] is the fully opaque color with sRGB components [r], [g],\n [b] in \\[0, 1\\], converted to OKLCH. *)\n\nval rgba : r:float -> g:float -> b:float -> a:float -> unit -> t\n(** [rgba ~r ~g ~b ~a ()] is like {!rgb} with alpha [a]. *)\n\nval hex : string -> t\n(** [hex s] is the color parsed from the hex string [s]. Accepts [\"#RRGGBB\"] and\n [\"#RRGGBBAA\"] formats.\n\n Raises [Invalid_argument] if [s] is not a valid hex color. *)\n\n(** {1:accessors Accessors} *)\n\nval lightness : t -> float\n(** [lightness c] is the OKLCH lightness of [c] in \\[0, 1\\]. *)\n\nval chroma : t -> float\n(** [chroma c] is the OKLCH chroma of [c] in \\[0, ~0.4\\]. *)\n\nval hue : t -> float\n(** [hue c] is the OKLCH hue of [c] in degrees \\[0, 360). *)\n\nval alpha : t -> float\n(** [alpha c] is the alpha of [c] in \\[0, 1\\]. *)\n\n(** {1:converting Converting} *)\n\nval to_rgba : t -> float * float * float * float\n(** [to_rgba c] is [(r, g, b, a)] with sRGB components in \\[0, 1\\]. Values are\n clamped to the sRGB gamut. *)\n\n(** {1:operations Operations} *)\n\nval with_alpha : float -> t -> t\n(** [with_alpha a c] is [c] with alpha set to [a]. *)\n\nval lighten : float -> t -> t\n(** [lighten amount c] is [c] with lightness increased by [amount], clamped to\n \\[0, 1\\]. *)\n\nval darken : float -> t -> t\n(** [darken amount c] is [c] with lightness decreased by [amount], clamped to\n \\[0, 1\\]. *)\n\nval mix : float -> t -> t -> t\n(** [mix ratio a b] is the perceptual blend of [a] and [b]. [ratio] is the\n interpolation factor: [0.0] gives [a], [1.0] gives [b]. Hue is interpolated\n along the shortest arc. *)\n\n(** {1:named Named colors}\n\n The default named colors follow the\n {{:https://jfly.uni-koeln.de/color/}Okabe-Ito} palette, designed to be\n distinguishable under all forms of color-vision deficiency. *)\n\nval orange : t\nval sky_blue : t\nval green : t\nval yellow : t\nval blue : t\nval vermillion : t\nval purple : t\nval black : t\nval white : t\nval gray : t\n\n(** {1:fmt Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp] formats the color as [oklch(L C H / A)]. *)\n"
},
{
"path": "packages/hugin/lib/dune",
"content": "(library\n (name hugin)\n (public_name hugin)\n (private_modules\n axis\n spec\n scale\n ticks\n scene\n prepared\n resolve\n image_util\n cairo_sdl\n cairo_backend\n svg_backend\n cmap_data)\n (libraries nx nx.buffer ucairo usdl))\n"
},
{
"path": "packages/hugin/lib/hugin.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Color = Color\nmodule Cmap = Cmap\nmodule Theme = Theme\n\ntype t = Spec.t\ntype marker = Spec.marker = Circle | Square | Triangle | Plus | Star\n\ntype legend_loc = Spec.legend_loc =\n | Upper_right\n | Upper_left\n | Lower_right\n | Lower_left\n | Center\n | Right\n | Upper_center\n | Lower_center\n\ntype line_style = Spec.line_style\ntype scale = Spec.scale\ntype stretch = Spec.stretch\n\n(* Mark constructors *)\n\nlet line = Spec.line\nlet point = Spec.point\nlet bar = Spec.bar\nlet hist = Spec.hist\nlet image = Spec.image\nlet text = Spec.text\nlet hline = Spec.hline\nlet vline = Spec.vline\nlet abline = Spec.abline\nlet fill_between = Spec.fill_between\nlet hspan = Spec.hspan\nlet vspan = Spec.vspan\nlet errorbar = Spec.errorbar\nlet heatmap = Spec.heatmap\nlet imshow = Spec.imshow\nlet contour = Spec.contour\n\n(* Composition *)\n\nlet layers = Spec.layers\n\n(* Decorations *)\n\nlet title = Spec.title\nlet xlabel = Spec.xlabel\nlet ylabel = Spec.ylabel\nlet xlim = Spec.xlim\nlet ylim = Spec.ylim\nlet xscale = Spec.xscale\nlet yscale = Spec.yscale\nlet xinvert = Spec.xinvert\nlet yinvert = Spec.yinvert\nlet grid_lines = Spec.grid_lines\nlet legend = Spec.legend\nlet xticks = Spec.xticks\nlet yticks = Spec.yticks\nlet with_theme = Spec.with_theme\nlet xtick_format = Spec.xtick_format\nlet ytick_format = Spec.ytick_format\nlet frame = Spec.frame\nlet no_axes = Spec.no_axes\n\n(* Layout *)\n\nlet grid = Spec.grid_layout\nlet hstack ?gap specs = Spec.grid_layout ?gap [ specs ]\nlet vstack ?gap specs = Spec.grid_layout ?gap (List.map (fun s -> [ s ]) specs)\n\n(* Rendering *)\n\nlet default_width = 1600.\nlet default_height = 1200.\n\n(* Use Cairo text measurement for all backends for consistent layout *)\nlet resolve_with_cairo ~theme ~width ~height spec =\n let surface = Ucairo.Image.create ~w:1 ~h:1 in\n let cr = Ucairo.create surface in\n let tm = Cairo_backend.text_measurer cr in\n let scene = Resolve.resolve ~text_measurer:tm ~theme ~width ~height spec in\n Ucairo.Surface.finish surface;\n scene\n\nlet show ?(theme = Theme.default) ?(width = default_width)\n ?(height = default_height) spec =\n let prepared = Prepared.compile ~theme spec in\n Cairo_backend.show_interactive ~theme ~width ~height prepared\n\nlet render_png ?(theme = Theme.default) ?(width = default_width)\n ?(height = default_height) filename spec =\n let scene = resolve_with_cairo ~theme ~width ~height spec in\n Cairo_backend.render_to_png filename ~width ~height scene\n\nlet render_pdf ?(theme = Theme.default) ?(width = default_width)\n ?(height = default_height) filename spec =\n let scene = resolve_with_cairo ~theme ~width ~height spec in\n Cairo_backend.render_to_pdf filename ~width ~height scene\n\nlet render_svg ?(theme = Theme.default) ?(width = default_width)\n ?(height = default_height) filename spec =\n let scene = resolve_with_cairo ~theme ~width ~height spec in\n Svg_backend.render_to_file filename scene\n\nlet render_svg_to_string ?(theme = Theme.default) ?(width = default_width)\n ?(height = default_height) spec =\n let scene = resolve_with_cairo ~theme ~width ~height spec in\n Svg_backend.render scene\n\nlet render_to_buffer ?(theme = Theme.default) ?(width = default_width)\n ?(height = default_height) spec =\n let scene = resolve_with_cairo ~theme ~width ~height spec in\n Cairo_backend.render_to_buffer ~width ~height scene\n\nlet infer_dimensions spec =\n let rec grid_shape = function\n | Spec.Grid { rows; _ } ->\n let nrows = List.length rows in\n let ncols =\n List.fold_left (fun acc row -> max acc (List.length row)) 0 rows\n in\n Some (nrows, ncols)\n | Spec.Decorated { inner; _ } -> grid_shape inner\n | _ -> None\n in\n match grid_shape spec with\n | Some (nrows, ncols) when ncols > 0 ->\n let cell_w = default_width /. float ncols in\n (default_width, cell_w *. float nrows)\n | _ -> (default_width, default_height)\n\nlet pp fmt spec =\n let width, height = infer_dimensions spec in\n let buf = render_to_buffer ~width ~height spec in\n let b64 = Image_util.base64_encode buf in\n Format.fprintf fmt \"\" b64\n"
},
{
"path": "packages/hugin/lib/hugin.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Declarative plotting and visualization.\n\n Hugin turns immutable plot specifications into rendered output. A plot is a\n value of type {!t} built from mark constructors ({!line}, {!point}, {!bar},\n {!hist}), composed with {!layers}, decorated with {!title}, {!xlabel}, etc.\n via the [|>] pipeline, and rendered with {!show}, {!render_png}, or\n {!render_svg}.\n\n {[\n let x = Nx.linspace Float32 0. 6.28 100 in\n let y = Nx.map (fun v -> Float.sin v) x in\n Hugin.line ~x ~y () |> Hugin.title \"Sine wave\"\n |> Hugin.render_png \"sine.png\"\n ]} *)\n\n(** {1:sub Sub-modules} *)\n\nmodule Color = Color\n(** Perceptually uniform OKLCH colors. *)\n\nmodule Cmap = Cmap\n(** Colormaps. *)\n\nmodule Theme = Theme\n(** Visual themes. *)\n\n(** {1:types Types} *)\n\ntype t\n(** The type for plot specifications. Immutable and composable. *)\n\ntype marker =\n | Circle\n | Square\n | Triangle\n | Plus\n | Star (** The type for point marker shapes. *)\n\ntype legend_loc =\n | Upper_right\n | Upper_left\n | Lower_right\n | Lower_left\n | Center\n | Right\n | Upper_center\n | Lower_center (** The type for legend placement. *)\n\ntype line_style = [ `Solid | `Dashed | `Dotted | `Dash_dot ]\n(** The type for line dash patterns. *)\n\ntype scale = [ `Linear | `Log | `Sqrt | `Asinh | `Symlog of float ]\n(** The type for axis scales. [`Sqrt] and [`Asinh] handle zero gracefully.\n [`Symlog linthresh] is linear within \\[[-linthresh];[linthresh]\\] and\n logarithmic outside. *)\n\ntype stretch = [ `Linear | `Log | `Sqrt | `Asinh | `Power of float ]\n(** The type for image stretch functions. [`Power a] raises normalized values to\n the power [a]. *)\n\n(** {1:marks Mark constructors}\n\n Each constructor builds a single-layer specification from data arrays and\n optional visual properties. A mark is already a valid {!t} that can be\n rendered directly. *)\n\nval line :\n x:Nx.float32_t ->\n y:Nx.float32_t ->\n ?color:Color.t ->\n ?line_width:float ->\n ?line_style:line_style ->\n ?step:[ `Pre | `Post | `Mid ] ->\n ?marker:marker ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [line ~x ~y ()] is a line plot connecting the points [(x.(i), y.(i))].\n\n [color] defaults to the next color in the theme palette. [line_width]\n defaults to the theme line width. [line_style] defaults to [`Solid]. [step]\n draws a staircase line: [`Post] holds each value until the next x-point,\n [`Pre] steps to the new value at the current x-point, [`Mid] steps at the\n midpoint between consecutive x-points. *)\n\nval point :\n x:Nx.float32_t ->\n y:Nx.float32_t ->\n ?color:Color.t ->\n ?color_by:Nx.float32_t ->\n ?size:float ->\n ?size_by:Nx.float32_t ->\n ?marker:marker ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [point ~x ~y ()] is a scatter plot of discrete markers at [(x.(i), y.(i))].\n\n [color_by] maps per-point values through the theme's sequential colormap.\n [size_by] scales marker area per point. [marker] defaults to {!Circle}. *)\n\nval bar :\n x:Nx.float32_t ->\n height:Nx.float32_t ->\n ?width:float ->\n ?bottom:float ->\n ?color:Color.t ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [bar ~x ~height ()] is a bar chart with bars centered on [x] values,\n extending from [bottom] (default [0.0]) to [bottom + height]. [width]\n defaults to [0.8]. *)\n\nval hist :\n x:Nx.float32_t ->\n ?bins:[ `Num of int | `Edges of float array ] ->\n ?density:bool ->\n ?color:Color.t ->\n ?label:string ->\n unit ->\n t\n(** [hist ~x ()] is a histogram of the values in [x].\n\n [bins] defaults to [`Num 10]. When [density] is [true], the histogram is\n normalized so the total area equals [1.0]. *)\n\nval image : ?extent:float * float * float * float -> Nx.uint8_t -> t\n(** [image ?extent data] displays [data] as an image. [data] has shape\n [[|h; w; 3|]] (RGB) or [[|h; w; 4|]] (RGBA).\n\n When [extent] is [(xmin, xmax, ymin, ymax)], the image is placed in data\n coordinates. Without [extent], the image is centered in the plot area\n preserving aspect ratio. *)\n\nval text :\n x:float ->\n y:float ->\n string ->\n ?color:Color.t ->\n ?font_size:float ->\n unit ->\n t\n(** [text ~x ~y s ()] places the string [s] at data coordinates [(x, y)]. *)\n\nval hline :\n y:float ->\n ?color:Color.t ->\n ?line_width:float ->\n ?line_style:line_style ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [hline ~y ()] draws a horizontal reference line at [y] spanning the full\n plot width. *)\n\nval vline :\n x:float ->\n ?color:Color.t ->\n ?line_width:float ->\n ?line_style:line_style ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [vline ~x ()] draws a vertical reference line at [x] spanning the full plot\n height. *)\n\nval abline :\n slope:float ->\n intercept:float ->\n ?color:Color.t ->\n ?line_width:float ->\n ?line_style:line_style ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [abline ~slope ~intercept ()] draws a diagonal line\n [y = slope * x + intercept] spanning the full plot area. Useful for\n regression lines and [y = x] references. *)\n\nval fill_between :\n x:Nx.float32_t ->\n y1:Nx.float32_t ->\n y2:Nx.float32_t ->\n ?where:Nx.float32_t ->\n ?color:Color.t ->\n ?alpha:float ->\n ?label:string ->\n unit ->\n t\n(** [fill_between ~x ~y1 ~y2 ()] fills the area between curves [y1] and [y2]\n over the shared [x] axis. [alpha] defaults to [0.3].\n\n [where] is an optional mask array of the same length as [x]: the fill is\n only drawn where [where.(i) > 0.], producing separate filled regions. *)\n\nval hspan :\n y0:float ->\n y1:float ->\n ?color:Color.t ->\n ?alpha:float ->\n ?label:string ->\n unit ->\n t\n(** [hspan ~y0 ~y1 ()] is a horizontal shaded band between [y0] and [y1],\n spanning the full plot width. [alpha] defaults to [0.2]. *)\n\nval vspan :\n x0:float ->\n x1:float ->\n ?color:Color.t ->\n ?alpha:float ->\n ?label:string ->\n unit ->\n t\n(** [vspan ~x0 ~x1 ()] is a vertical shaded band between [x0] and [x1], spanning\n the full plot height. [alpha] defaults to [0.2]. *)\n\nval errorbar :\n x:Nx.float32_t ->\n y:Nx.float32_t ->\n yerr:\n [ `Symmetric of Nx.float32_t | `Asymmetric of Nx.float32_t * Nx.float32_t ] ->\n ?xerr:\n [ `Symmetric of Nx.float32_t | `Asymmetric of Nx.float32_t * Nx.float32_t ] ->\n ?color:Color.t ->\n ?line_width:float ->\n ?cap_size:float ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [errorbar ~x ~y ~yerr ()] draws error bars at [(x.(i), y.(i))].\n\n [yerr] specifies vertical error: [`Symmetric e] draws [y +/- e],\n [`Asymmetric (lo, hi)] draws [[y - lo, y + hi]]. [xerr] adds horizontal\n error bars. [cap_size] defaults to half the theme marker size. *)\n\nval heatmap :\n data:Nx.float32_t ->\n ?annotate:bool ->\n ?cmap:Cmap.t ->\n ?vmin:float ->\n ?vmax:float ->\n ?fmt:(float -> string) ->\n unit ->\n t\n(** [heatmap ~data ()] displays a 2D array as a grid of colored cells. [data]\n has shape [[|rows; cols|]]. Row 0 appears at the top.\n\n [cmap] defaults to the theme's sequential colormap. [vmin] and [vmax]\n override the automatic value range. When [annotate] is [true], each cell\n shows its value formatted by [fmt] (default [Printf.sprintf \"%.2g\"]). *)\n\nval imshow :\n data:Nx.float32_t ->\n ?stretch:stretch ->\n ?cmap:Cmap.t ->\n ?vmin:float ->\n ?vmax:float ->\n unit ->\n t\n(** [imshow ~data ()] displays a 2D float array as a colormapped image. [data]\n has shape [[|rows; cols|]].\n\n [stretch] controls the transfer function applied before colormap lookup:\n [`Linear] (default), [`Log], [`Sqrt], [`Asinh], or [`Power a]. [cmap]\n defaults to the theme's sequential colormap. [vmin] and [vmax] override the\n automatic value range. *)\n\nval contour :\n data:Nx.float32_t ->\n x0:float ->\n x1:float ->\n y0:float ->\n y1:float ->\n ?levels:[ `Num of int | `Values of float array ] ->\n ?filled:bool ->\n ?cmap:Cmap.t ->\n ?color:Color.t ->\n ?line_width:float ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [contour ~data ~x0 ~x1 ~y0 ~y1 ()] draws iso-level contour lines through the\n 2D grid [data] of shape [[|rows; cols|]], mapped to the data-space rectangle\n \\[[x0];[x1]\\] x \\[[y0];[y1]\\].\n\n [levels] defaults to [`Num 8]. When [filled] is [true], regions between\n adjacent levels are filled. [color] sets a single stroke color for unfilled\n contours; [cmap] assigns per-level colors from the theme's sequential\n colormap. *)\n\n(** {1:composition Composition} *)\n\nval layers : t list -> t\n(** [layers marks] overlays [marks] on shared axes. A single mark is already a\n valid {!t}; [layers] is only needed to combine multiple marks into one plot.\n*)\n\n(** {1:decorations Decorations}\n\n Decoration functions add metadata to a specification. They are designed for\n the [|>] pipeline:\n {[\n line ~x ~y () |> title \"My Plot\" |> xlabel \"Time\"\n ]} *)\n\nval title : string -> t -> t\n(** [title s t] is [t] with plot title [s]. *)\n\nval xlabel : string -> t -> t\n(** [xlabel s t] is [t] with x-axis label [s]. *)\n\nval ylabel : string -> t -> t\n(** [ylabel s t] is [t] with y-axis label [s]. *)\n\nval xlim : float -> float -> t -> t\n(** [xlim lo hi t] is [t] with x-axis range fixed to \\[[lo];[hi]\\]. *)\n\nval ylim : float -> float -> t -> t\n(** [ylim lo hi t] is [t] with y-axis range fixed to \\[[lo];[hi]\\]. *)\n\nval xscale : scale -> t -> t\n(** [xscale s t] is [t] with x-axis scale [s]. Defaults to [`Linear].\n\n [`Sqrt] and [`Asinh] handle zero gracefully. [`Symlog linthresh] is linear\n within \\[[-linthresh];[linthresh]\\] and logarithmic outside. *)\n\nval yscale : scale -> t -> t\n(** [yscale s t] is [t] with y-axis scale [s]. Defaults to [`Linear]. *)\n\nval xinvert : t -> t\n(** [xinvert t] is [t] with the x-axis inverted (values increase right-to-left).\n Useful for right ascension in sky charts. *)\n\nval yinvert : t -> t\n(** [yinvert t] is [t] with the y-axis inverted (values increase top-to-bottom).\n Useful for magnitude axes in HR diagrams. *)\n\nval grid_lines : bool -> t -> t\n(** [grid_lines visible t] is [t] with grid lines shown or hidden. *)\n\nval legend : ?loc:legend_loc -> ?ncol:int -> t -> t\n(** [legend ?loc ?ncol t] is [t] with the legend shown at [loc]. [loc] defaults\n to {!Upper_right}. [ncol] defaults to [1]; set higher for multi-column\n layouts with many series. The legend is automatically visible when any mark\n has a [~label]. *)\n\nval xticks : (float * string) list -> t -> t\n(** [xticks ticks t] is [t] with explicit x-axis tick positions and labels.\n Overrides auto-generated ticks. *)\n\nval yticks : (float * string) list -> t -> t\n(** [yticks ticks t] is [t] with explicit y-axis tick positions and labels.\n Overrides auto-generated ticks. *)\n\nval with_theme : Theme.t -> t -> t\n(** [with_theme th t] is [t] rendered with theme [th] instead of the default. *)\n\nval xtick_format : (float -> string) -> t -> t\n(** [xtick_format fmt t] is [t] with x-axis tick labels formatted by [fmt].\n Overrides auto-generated labels while preserving tick positions. *)\n\nval ytick_format : (float -> string) -> t -> t\n(** [ytick_format fmt t] is [t] with y-axis tick labels formatted by [fmt].\n Overrides auto-generated labels while preserving tick positions. *)\n\nval frame : bool -> t -> t\n(** [frame visible t] is [t] with the axis border rectangle shown or hidden.\n [visible] defaults to [true]. *)\n\nval no_axes : t -> t\n(** [no_axes t] hides the axis frame, ticks, and tick labels. Title is\n preserved. The full panel area is used for marks. Useful for image grids:\n\n {[\n List.init 10 (fun i ->\n Hugin.imshow ~data:digits.(i) ~cmap:Cmap.gray ()\n |> Hugin.title (string_of_int labels.(i))\n |> Hugin.no_axes)\n |> Hugin.hstack\n ]} *)\n\n(** {1:layout Layout} *)\n\nval grid : ?gap:float -> t list list -> t\n(** [grid rows] arranges specifications in a grid. Each inner list is a row of\n panels. [gap] defaults to [0.05] (fraction of total size). *)\n\nval hstack : ?gap:float -> t list -> t\n(** [hstack specs] arranges [specs] in a single row. *)\n\nval vstack : ?gap:float -> t list -> t\n(** [vstack specs] arranges [specs] in a single column. *)\n\n(** {1:rendering Rendering} *)\n\nval show : ?theme:Theme.t -> ?width:float -> ?height:float -> t -> unit\n(** [show t] displays [t] in an interactive SDL window.\n\n [width] defaults to [1600.0]. [height] defaults to [1200.0]. The window\n supports resize (re-resolves at new dimensions) and closes on Escape or Q.\n*)\n\nval render_png :\n ?theme:Theme.t -> ?width:float -> ?height:float -> string -> t -> unit\n(** [render_png filename t] writes [t] as a PNG image to [filename].\n\n [width] defaults to [1600.0]. [height] defaults to [1200.0]. *)\n\nval render_pdf :\n ?theme:Theme.t -> ?width:float -> ?height:float -> string -> t -> unit\n(** [render_pdf filename t] writes [t] as a PDF document to [filename].\n\n [width] defaults to [1600.0]. [height] defaults to [1200.0]. *)\n\nval render_svg :\n ?theme:Theme.t -> ?width:float -> ?height:float -> string -> t -> unit\n(** [render_svg filename t] writes [t] as an SVG document to [filename].\n\n [width] defaults to [1600.0]. [height] defaults to [1200.0]. *)\n\nval render_svg_to_string :\n ?theme:Theme.t -> ?width:float -> ?height:float -> t -> string\n(** [render_svg_to_string t] is [t] rendered as an SVG document string.\n\n [width] defaults to [1600.0]. [height] defaults to [1200.0]. *)\n\nval render_to_buffer :\n ?theme:Theme.t -> ?width:float -> ?height:float -> t -> string\n(** [render_to_buffer t] is [t] rendered as a PNG image, returned as a string of\n bytes. *)\n\n(** {1:fmt Formatting} *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp] renders the specification as a PNG data URI. Intended for use with\n [#install_printer] in the toplevel and Quill.\n\n Output format: [] *)\n"
},
{
"path": "packages/hugin/lib/image_util.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Shared image encoding utilities *)\n\n(* Base64 *)\n\nlet base64_alphabet =\n \"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/\"\n\nlet base64_encode input =\n let len = String.length input in\n let out_len = (len + 2) / 3 * 4 in\n let out = Bytes.create out_len in\n let rec loop i j =\n if i < len then begin\n let b0 = Char.code (String.unsafe_get input i) in\n let b1 =\n if i + 1 < len then Char.code (String.unsafe_get input (i + 1)) else 0\n in\n let b2 =\n if i + 2 < len then Char.code (String.unsafe_get input (i + 2)) else 0\n in\n Bytes.unsafe_set out j (String.unsafe_get base64_alphabet (b0 lsr 2));\n Bytes.unsafe_set out (j + 1)\n (String.unsafe_get base64_alphabet (((b0 land 3) lsl 4) lor (b1 lsr 4)));\n Bytes.unsafe_set out (j + 2)\n (if i + 1 < len then\n String.unsafe_get base64_alphabet\n (((b1 land 0xf) lsl 2) lor (b2 lsr 6))\n else '=');\n Bytes.unsafe_set out (j + 3)\n (if i + 2 < len then String.unsafe_get base64_alphabet (b2 land 0x3f)\n else '=');\n loop (i + 3) (j + 4)\n end\n in\n loop 0 0;\n Bytes.unsafe_to_string out\n\n(* Nx uint8 image -> Cairo ARGB32 surface *)\n\nlet nx_to_cairo_surface (data : Nx.uint8_t) =\n let shape = Nx.shape data in\n let img_h = shape.(0) and img_w = shape.(1) in\n let channels = if Array.length shape > 2 then shape.(2) else 1 in\n let stride = Ucairo.Image.stride_for_width img_w in\n let data_arr =\n Bigarray.Array1.create Bigarray.int8_unsigned Bigarray.c_layout\n (stride * img_h)\n in\n let buf = Nx.data data in\n let base = Nx.offset data in\n let strides = Nx.strides data in\n (* uint8: byte strides = element strides *)\n let s0 = strides.(0) and s1 = strides.(1) in\n let s2 = if Array.length strides > 2 then strides.(2) else 0 in\n for row = 0 to img_h - 1 do\n let row_base = base + (row * s0) in\n for col = 0 to img_w - 1 do\n let off = (row * stride) + (col * 4) in\n let idx = row_base + (col * s1) in\n let r = Nx_buffer.unsafe_get buf idx in\n let g = Nx_buffer.unsafe_get buf (idx + s2) in\n let b = Nx_buffer.unsafe_get buf (idx + (2 * s2)) in\n let a =\n if channels >= 4 then Nx_buffer.unsafe_get buf (idx + (3 * s2)) else 255\n in\n (* Cairo ARGB32: premultiplied BGRA in memory on little-endian *)\n let premul c a = c * a / 255 in\n Bigarray.Array1.unsafe_set data_arr off (premul b a);\n Bigarray.Array1.unsafe_set data_arr (off + 1) (premul g a);\n Bigarray.Array1.unsafe_set data_arr (off + 2) (premul r a);\n Bigarray.Array1.unsafe_set data_arr (off + 3) a\n done\n done;\n Ucairo.Image.create_for_data8 data_arr ~w:img_w ~h:img_h ~stride\n\nlet nx_to_png_base64 data =\n let surface = nx_to_cairo_surface data in\n let png_buf = Buffer.create 4096 in\n Ucairo.Png.write_to_stream surface (Buffer.add_string png_buf);\n Ucairo.Surface.finish surface;\n base64_encode (Buffer.contents png_buf)\n"
},
{
"path": "packages/hugin/lib/image_util.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Image encoding utilities.\n\n {b Internal module.} Shared base64 encoding and Nx-to-Cairo-surface\n conversion used by both rendering backends. *)\n\n(** {1:base64 Base64} *)\n\nval base64_encode : string -> string\n(** [base64_encode s] is the base64 encoding of [s]. *)\n\n(** {1:surface Cairo surface conversion} *)\n\nval nx_to_cairo_surface : Nx.uint8_t -> Ucairo.surface\n(** [nx_to_cairo_surface data] is a Cairo ARGB32 image surface from [data].\n [data] has shape [[|h; w; 3|]] (RGB) or [[|h; w; 4|]] (RGBA). *)\n\nval nx_to_png_base64 : Nx.uint8_t -> string\n(** [nx_to_png_base64 data] is the base64-encoded PNG of [data]. *)\n"
},
{
"path": "packages/hugin/lib/prepared.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Data-only compilation: Spec.t -> Prepared.t\n\n Compiles once per dataset. Separates data-dependent work (collecting\n decorations, histogram binning, auto-coloring, data bounds) from\n layout-dependent work (pixel coordinates, text measurement) which lives in\n Resolve. *)\n\n(* Data bounds *)\n\nlet nx_finite_range (arr : Nx.float32_t) =\n let n = (Nx.shape arr).(0) in\n let lo = ref Float.infinity and hi = ref Float.neg_infinity in\n for i = 0 to n - 1 do\n let v = Nx.item [ i ] arr in\n if Float.is_finite v then begin\n if v < !lo then lo := v;\n if v > !hi then hi := v\n end\n done;\n (!lo, !hi)\n\nlet expand_range scale lo hi =\n if lo = hi then (lo -. 1., hi +. 1.)\n else\n match scale with\n | `Log ->\n let lo_log = Float.log10 (Float.max 1e-10 lo) in\n let hi_log = Float.log10 (Float.max 1e-10 hi) in\n let pad = (hi_log -. lo_log) *. 0.05 in\n (Float.pow 10. (lo_log -. pad), Float.pow 10. (hi_log +. pad))\n | `Sqrt ->\n let lo = Float.max 0. lo in\n let pad = (hi -. lo) *. 0.05 in\n (Float.max 0. (lo -. pad), hi +. pad)\n | `Asinh | `Symlog _ | `Linear ->\n let pad = (hi -. lo) *. 0.05 in\n (lo -. pad, hi +. pad)\n\nlet mark_x_range = function\n | Spec.Line { x; _ } | Spec.Point { x; _ } -> Some (nx_finite_range x)\n | Spec.Bar { x; width; _ } ->\n let lo, hi = nx_finite_range x in\n let w = (match width with Some w -> w | None -> 0.8) /. 2. in\n Some (lo -. w, hi +. w)\n | Spec.Hist { x; _ } -> Some (nx_finite_range x)\n | Spec.Image { extent = Some (xmin, xmax, _, _); _ } ->\n Some (Float.min xmin xmax, Float.max xmin xmax)\n | Spec.Image _ -> None\n | Spec.Text_mark { x; _ } -> Some (x, x)\n | Spec.Hline _ -> None\n | Spec.Vline { x; _ } -> Some (x, x)\n | Spec.Abline _ -> None\n | Spec.Fill_between { x; _ } -> Some (nx_finite_range x)\n | Spec.Errorbar { x; xerr; _ } ->\n let lo, hi = nx_finite_range x in\n let lo, hi =\n match xerr with\n | Some (`Symmetric e) ->\n let _, emax = nx_finite_range e in\n (lo -. emax, hi +. emax)\n | Some (`Asymmetric (elo, ehi)) ->\n let _, emlo = nx_finite_range elo in\n let _, emhi = nx_finite_range ehi in\n (lo -. emlo, hi +. emhi)\n | None -> (lo, hi)\n in\n Some (lo, hi)\n | Spec.Hspan _ -> None\n | Spec.Vspan { x0; x1; _ } -> Some (Float.min x0 x1, Float.max x0 x1)\n | Spec.Heatmap { data; _ } ->\n let shape = Nx.shape data in\n let cols = float shape.(1) in\n Some (0., cols)\n | Spec.Imshow _ -> None\n | Spec.Contour { x0; x1; _ } -> Some (Float.min x0 x1, Float.max x0 x1)\n\nlet mark_y_range = function\n | Spec.Line { y; _ } | Spec.Point { y; _ } -> Some (nx_finite_range y)\n | Spec.Bar { height; bottom; _ } ->\n let lo, hi = nx_finite_range height in\n Some (Float.min bottom (bottom +. lo), Float.max bottom (bottom +. hi))\n | Spec.Hist _ -> None\n | Spec.Image { extent = Some (_, _, ymin, ymax); _ } ->\n Some (Float.min ymin ymax, Float.max ymin ymax)\n | Spec.Image _ -> None\n | Spec.Text_mark { y; _ } -> Some (y, y)\n | Spec.Hline { y; _ } -> Some (y, y)\n | Spec.Vline _ -> None\n | Spec.Abline _ -> None\n | Spec.Fill_between { y1; y2; _ } ->\n let lo1, hi1 = nx_finite_range y1 in\n let lo2, hi2 = nx_finite_range y2 in\n Some (Float.min lo1 lo2, Float.max hi1 hi2)\n | Spec.Errorbar { y; yerr; _ } ->\n let lo, hi = nx_finite_range y in\n let lo, hi =\n match yerr with\n | `Symmetric e ->\n let _, emax = nx_finite_range e in\n (lo -. emax, hi +. emax)\n | `Asymmetric (elo, ehi) ->\n let _, emlo = nx_finite_range elo in\n let _, emhi = nx_finite_range ehi in\n (lo -. emlo, hi +. emhi)\n in\n Some (lo, hi)\n | Spec.Hspan { y0; y1; _ } -> Some (Float.min y0 y1, Float.max y0 y1)\n | Spec.Vspan _ -> None\n | Spec.Heatmap { data; _ } ->\n let shape = Nx.shape data in\n let rows = float shape.(0) in\n Some (0., rows)\n | Spec.Imshow _ -> None\n | Spec.Contour { y0; y1; _ } -> Some (Float.min y0 y1, Float.max y0 y1)\n\nlet union_range a b =\n match (a, b) with\n | None, x | x, None -> x\n | Some (a0, a1), Some (b0, b1) -> Some (Float.min a0 b0, Float.max a1 b1)\n\nlet compute_data_bounds ~xscale ~yscale marks =\n let xr =\n List.fold_left (fun acc m -> union_range acc (mark_x_range m)) None marks\n in\n let yr =\n List.fold_left (fun acc m -> union_range acc (mark_y_range m)) None marks\n in\n let xlo, xhi =\n match xr with Some (a, b) -> expand_range xscale a b | None -> (0., 1.)\n in\n let ylo, yhi =\n match yr with Some (a, b) -> expand_range yscale a b | None -> (0., 1.)\n in\n (xlo, xhi, ylo, yhi)\n\n(* Collect decorations from spec tree *)\n\ntype collected = {\n marks : Spec.mark list;\n x : Axis.config;\n y : Axis.config;\n title : string option;\n grid_visible : bool option;\n frame_visible : bool option;\n legend_loc : Spec.legend_loc option;\n legend_ncol : int;\n theme_override : Theme.t option;\n}\n\nlet empty_collected =\n {\n marks = [];\n x = Axis.empty_config;\n y = Axis.empty_config;\n title = None;\n grid_visible = None;\n frame_visible = None;\n legend_loc = None;\n legend_ncol = 1;\n theme_override = None;\n }\n\nlet rec collect c = function\n | Spec.Mark m -> { c with marks = m :: c.marks }\n | Spec.Layers ts -> List.fold_left collect c ts\n | Spec.Decorated { inner; decorations } ->\n let c = collect c inner in\n List.fold_left apply_decoration c decorations\n | Spec.Grid _ -> c\n\nand apply_decoration c = function\n | Spec.Title s when c.title = None -> { c with title = Some s }\n | Spec.Xlabel s when c.x.label = None ->\n { c with x = { c.x with label = Some s } }\n | Spec.Ylabel s when c.y.label = None ->\n { c with y = { c.y with label = Some s } }\n | Spec.Xlim (lo, hi) when c.x.lim = None ->\n { c with x = { c.x with lim = Some (lo, hi) } }\n | Spec.Ylim (lo, hi) when c.y.lim = None ->\n { c with y = { c.y with lim = Some (lo, hi) } }\n | Spec.Xscale s when c.x.scale = None ->\n { c with x = { c.x with scale = Some s } }\n | Spec.Yscale s when c.y.scale = None ->\n { c with y = { c.y with scale = Some s } }\n | Spec.Xinvert -> { c with x = { c.x with invert = true } }\n | Spec.Yinvert -> { c with y = { c.y with invert = true } }\n | Spec.Grid_visible v when c.grid_visible = None ->\n { c with grid_visible = Some v }\n | Spec.Legend (loc, ncol) when c.legend_loc = None ->\n { c with legend_loc = Some loc; legend_ncol = ncol }\n | Spec.Xticks t when c.x.ticks = None ->\n { c with x = { c.x with ticks = Some t } }\n | Spec.Yticks t when c.y.ticks = None ->\n { c with y = { c.y with ticks = Some t } }\n | Spec.With_theme t when c.theme_override = None ->\n { c with theme_override = Some t }\n | Spec.Xtick_format f when c.x.tick_format = None ->\n { c with x = { c.x with tick_format = Some f } }\n | Spec.Ytick_format f when c.y.tick_format = None ->\n { c with y = { c.y with tick_format = Some f } }\n | Spec.Frame v when c.frame_visible = None ->\n { c with frame_visible = Some v }\n | _ -> c\n\n(* Auto-coloring *)\n\nlet mark_color = function\n | Spec.Line { color; _ }\n | Spec.Point { color; _ }\n | Spec.Bar { color; _ }\n | Spec.Hist { color; _ }\n | Spec.Text_mark { color; _ }\n | Spec.Hline { color; _ }\n | Spec.Vline { color; _ }\n | Spec.Abline { color; _ }\n | Spec.Fill_between { color; _ }\n | Spec.Hspan { color; _ }\n | Spec.Vspan { color; _ }\n | Spec.Errorbar { color; _ }\n | Spec.Contour { color; _ } ->\n color\n | Spec.Image _ | Spec.Heatmap _ | Spec.Imshow _ -> None\n\nlet auto_color (theme : Theme.t) marks =\n let n_palette = Array.length theme.palette in\n List.mapi\n (fun i m ->\n match mark_color m with\n | Some _ -> m\n | None -> (\n let c = theme.palette.(i mod n_palette) in\n match m with\n | Spec.Line r -> Spec.Line { r with color = Some c }\n | Spec.Point r -> Spec.Point { r with color = Some c }\n | Spec.Bar r -> Spec.Bar { r with color = Some c }\n | Spec.Hist r -> Spec.Hist { r with color = Some c }\n | Spec.Hline r -> Spec.Hline { r with color = Some c }\n | Spec.Vline r -> Spec.Vline { r with color = Some c }\n | Spec.Abline r -> Spec.Abline { r with color = Some c }\n | Spec.Fill_between r -> Spec.Fill_between { r with color = Some c }\n | Spec.Hspan r -> Spec.Hspan { r with color = Some c }\n | Spec.Vspan r -> Spec.Vspan { r with color = Some c }\n | Spec.Errorbar r -> Spec.Errorbar { r with color = Some c }\n | Spec.Contour r -> Spec.Contour { r with color = Some c }\n | m -> m))\n marks\n\n(* Histogram normalization — convert Hist to Bar *)\n\nlet normalize_hist marks =\n List.map\n (fun m ->\n match m with\n | Spec.Hist { x; bins; density; color; label } ->\n let xmin, xmax = nx_finite_range x in\n let edges =\n match bins with\n | `Num num_bins ->\n Array.init (num_bins + 1) (fun i ->\n xmin +. ((xmax -. xmin) *. float i /. float num_bins))\n | `Edges e -> e\n in\n let num_bins = Array.length edges - 1 in\n let n = (Nx.shape x).(0) in\n let counts = Array.make num_bins 0. in\n let binned = ref 0 in\n for i = 0 to n - 1 do\n let v = Nx.item [ i ] x in\n if Float.is_finite v && v >= edges.(0) && v <= edges.(num_bins) then begin\n incr binned;\n let bin = ref 0 in\n while !bin < num_bins - 1 && v >= edges.(!bin + 1) do\n incr bin\n done;\n counts.(!bin) <- counts.(!bin) +. 1.\n end\n done;\n if density then begin\n let total =\n let b = float !binned in\n if b = 0. then 1. else b\n in\n for i = 0 to num_bins - 1 do\n let w = edges.(i + 1) -. edges.(i) in\n counts.(i) <- counts.(i) /. (total *. w)\n done\n end;\n let bar_x =\n Nx.init Float32 [| num_bins |] (fun idx ->\n let i = idx.(0) in\n (edges.(i) +. edges.(i + 1)) /. 2.)\n in\n let bar_h =\n Nx.init Float32 [| num_bins |] (fun idx -> counts.(idx.(0)))\n in\n let w = if num_bins > 0 then edges.(1) -. edges.(0) else 1. in\n Spec.Bar\n {\n x = bar_x;\n height = bar_h;\n width = Some w;\n bottom = 0.;\n color;\n label;\n alpha = None;\n }\n | m -> m)\n marks\n\n(* Guide ranges *)\n\nlet color_by_range marks =\n List.fold_left\n (fun acc m ->\n match m with\n | Spec.Point { color_by = Some cb; _ } ->\n let lo, hi = nx_finite_range cb in\n union_range acc (Some (lo, hi))\n | _ -> acc)\n None marks\n\nlet size_by_range marks =\n List.fold_left\n (fun acc m ->\n match m with\n | Spec.Point { size_by = Some sb; _ } ->\n let lo, hi = nx_finite_range sb in\n union_range acc (Some (lo, hi))\n | _ -> acc)\n None marks\n\n(* Collect marks from all panels in a spec tree *)\n\nlet rec collect_all_marks = function\n | Spec.Mark m -> [ m ]\n | Spec.Layers ts -> List.concat_map collect_all_marks ts\n | Spec.Decorated { inner; _ } -> collect_all_marks inner\n | Spec.Grid { rows; _ } ->\n List.concat_map (List.concat_map collect_all_marks) rows\n\n(* Grid-level decorations *)\n\ntype grid_decorations = {\n gd_title : string option;\n gd_xlabel : string option;\n gd_ylabel : string option;\n gd_legend_loc : Spec.legend_loc option;\n gd_legend_ncol : int;\n gd_theme_override : Theme.t option;\n}\n\nlet extract_grid_decorations decorations =\n let d =\n {\n gd_title = None;\n gd_xlabel = None;\n gd_ylabel = None;\n gd_legend_loc = None;\n gd_legend_ncol = 1;\n gd_theme_override = None;\n }\n in\n List.fold_left\n (fun d dec ->\n match dec with\n | Spec.Title s when d.gd_title = None -> { d with gd_title = Some s }\n | Spec.Xlabel s when d.gd_xlabel = None -> { d with gd_xlabel = Some s }\n | Spec.Ylabel s when d.gd_ylabel = None -> { d with gd_ylabel = Some s }\n | Spec.Legend (loc, ncol) when d.gd_legend_loc = None ->\n { d with gd_legend_loc = Some loc; gd_legend_ncol = ncol }\n | Spec.With_theme t when d.gd_theme_override = None ->\n { d with gd_theme_override = Some t }\n | _ -> d)\n d decorations\n\n(* Prepared panel — all data-only work done *)\n\ntype panel = {\n marks : Spec.mark list;\n x : Axis.t;\n y : Axis.t;\n title : string option;\n legend_loc : Spec.legend_loc option;\n legend_ncol : int;\n grid_visible : bool option;\n frame_visible : bool option;\n theme_override : Theme.t option;\n colorbar_range : (float * float) option;\n size_by_range : (float * float) option;\n}\n\ntype t =\n | Panel of panel\n | Grid of { rows : t list list; gap : float }\n | Decorated_grid of {\n decorations : grid_decorations;\n inner : t;\n all_marks : Spec.mark list;\n }\n\n(* Imshow: rasterize float32 data to uint8 RGB via stretch + colormap *)\n\nlet apply_stretch stretch v =\n match stretch with\n | `Linear -> v\n | `Log -> Float.log10 (1. +. (9. *. v)) /. Float.log10 10.\n | `Sqrt -> Float.sqrt (Float.max 0. v)\n | `Asinh ->\n let a = 10. in\n Float.asinh (a *. v) /. Float.asinh a\n | `Power a -> Float.pow (Float.max 0. v) a\n\nlet rasterize_imshow ~stretch ~cmap ~vmin ~vmax (data : Nx.float32_t) =\n let shape = Nx.shape data in\n let rows = shape.(0) and cols = shape.(1) in\n let lo = ref Float.infinity and hi = ref Float.neg_infinity in\n for r = 0 to rows - 1 do\n for c = 0 to cols - 1 do\n let v = Nx.item [ r; c ] data in\n if Float.is_finite v then begin\n if v < !lo then lo := v;\n if v > !hi then hi := v\n end\n done\n done;\n let vlo = match vmin with Some v -> v | None -> !lo in\n let vhi = match vmax with Some v -> v | None -> !hi in\n let vrange = if vhi = vlo then 1. else vhi -. vlo in\n let rgb = Nx.zeros Nx.uint8 [| rows; cols; 3 |] in\n for r = 0 to rows - 1 do\n for c = 0 to cols - 1 do\n let v = Nx.item [ r; c ] data in\n let t = Float.max 0. (Float.min 1. ((v -. vlo) /. vrange)) in\n let t = apply_stretch stretch t in\n let t = Float.max 0. (Float.min 1. t) in\n let color = Cmap.eval cmap t in\n let cr, cg, cb, _ = Color.to_rgba color in\n Nx.set_item [ r; c; 0 ] (int_of_float (cr *. 255.)) rgb;\n Nx.set_item [ r; c; 1 ] (int_of_float (cg *. 255.)) rgb;\n Nx.set_item [ r; c; 2 ] (int_of_float (cb *. 255.)) rgb\n done\n done;\n rgb\n\nlet normalize_imshow (theme : Theme.t) marks =\n List.map\n (fun m ->\n match m with\n | Spec.Imshow { data; stretch; cmap; vmin; vmax } ->\n let cmap = match cmap with Some c -> c | None -> theme.sequential in\n let rgb = rasterize_imshow ~stretch ~cmap ~vmin ~vmax data in\n Spec.Image { data = rgb; extent = None }\n | m -> m)\n marks\n\n(* Contour tracing via marching squares *)\n\ntype contour_paths = { level : float; paths : (float * float) array list }\n\n(* Join 2-point segments that share endpoints into connected polylines. Marching\n squares produces one segment per cell edge crossing. Segments from adjacent\n cells share exact floating-point endpoints (deterministic lerp), so we chain\n them with exact equality via a hashtable. *)\nlet join_segments segments =\n let n = List.length segments in\n if n = 0 then []\n else\n let segs = Array.of_list segments in\n let visited = Array.make n false in\n let adj = Hashtbl.create (2 * n) in\n Array.iteri\n (fun i (a, b) ->\n let add pt =\n let cur = try Hashtbl.find adj pt with Not_found -> [] in\n Hashtbl.replace adj pt (i :: cur)\n in\n add a;\n add b)\n segs;\n let find_unvisited_neighbor pt =\n match Hashtbl.find adj pt with\n | exception Not_found -> None\n | neighbors ->\n let rec scan = function\n | [] -> None\n | j :: rest -> if visited.(j) then scan rest else Some j\n in\n scan neighbors\n in\n let chains = ref [] in\n for start = 0 to n - 1 do\n if not visited.(start) then begin\n visited.(start) <- true;\n let a0, b0 = segs.(start) in\n (* front: backward extensions (cons'd, so in chain order). back: forward\n extensions (cons'd, so reversed). *)\n let front = ref [ a0 ] in\n let back = ref [ b0 ] in\n (* Extend forward from b0 *)\n let cur = ref b0 in\n let go = ref true in\n while !go do\n match find_unvisited_neighbor !cur with\n | None -> go := false\n | Some j ->\n visited.(j) <- true;\n let a, b = segs.(j) in\n let next = if a = !cur then b else a in\n back := next :: !back;\n cur := next\n done;\n (* Extend backward from a0 *)\n cur := a0;\n go := true;\n while !go do\n match find_unvisited_neighbor !cur with\n | None -> go := false\n | Some j ->\n visited.(j) <- true;\n let a, b = segs.(j) in\n let next = if a = !cur then b else a in\n front := next :: !front;\n cur := next\n done;\n (* front is in chain order; back is reversed *)\n chains := Array.of_list (!front @ List.rev !back) :: !chains\n end\n done;\n List.rev !chains\n\nlet trace_contours ~rows ~cols (data : Nx.float32_t) levels =\n let get r c =\n if r >= 0 && r < rows && c >= 0 && c < cols then Nx.item [ r; c ] data\n else 0.\n in\n List.map\n (fun level ->\n let segments = ref [] in\n for r = 0 to rows - 2 do\n for c = 0 to cols - 2 do\n let v00 = get r c in\n let v10 = get r (c + 1) in\n let v11 = get (r + 1) (c + 1) in\n let v01 = get (r + 1) c in\n let b0 = if v00 >= level then 1 else 0 in\n let b1 = if v10 >= level then 1 else 0 in\n let b2 = if v11 >= level then 1 else 0 in\n let b3 = if v01 >= level then 1 else 0 in\n let case = b0 lor (b1 lsl 1) lor (b2 lsl 2) lor (b3 lsl 3) in\n let lerp va vb =\n let d = vb -. va in\n if Float.abs d < 1e-30 then 0.5 else (level -. va) /. d\n in\n let fc = float c and fr = float r in\n let top = (fc +. lerp v00 v10, fr) in\n let right = (fc +. 1., fr +. lerp v10 v11) in\n let bottom = (fc +. lerp v01 v11, fr +. 1.) in\n let left = (fc, fr +. lerp v00 v01) in\n let add a b = segments := (a, b) :: !segments in\n begin match case with\n | 0 | 15 -> ()\n | 1 | 14 -> add top left\n | 2 | 13 -> add top right\n | 3 | 12 -> add left right\n | 4 | 11 -> add right bottom\n | 5 ->\n let center = (v00 +. v10 +. v11 +. v01) /. 4. in\n if center >= level then begin\n add top right;\n add bottom left\n end\n else begin\n add top left;\n add bottom right\n end\n | 6 | 9 -> add top bottom\n | 7 | 8 -> add bottom left\n | 10 ->\n let center = (v00 +. v10 +. v11 +. v01) /. 4. in\n if center >= level then begin\n add top left;\n add bottom right\n end\n else begin\n add top right;\n add bottom left\n end\n | _ -> ()\n end\n done\n done;\n let paths = join_segments !segments in\n { level; paths })\n levels\n\nlet prepare_contour ~x0 ~x1 ~y0 ~y1 ~data ~levels =\n let shape = Nx.shape data in\n let rows = shape.(0) and cols = shape.(1) in\n let lo = ref Float.infinity and hi = ref Float.neg_infinity in\n for r = 0 to rows - 1 do\n for c = 0 to cols - 1 do\n let v = Nx.item [ r; c ] data in\n if Float.is_finite v then begin\n if v < !lo then lo := v;\n if v > !hi then hi := v\n end\n done\n done;\n let vlo = !lo and vhi = !hi in\n let level_values =\n match levels with\n | `Values a -> Array.to_list a\n | `Num n ->\n let range = vhi -. vlo in\n if range = 0. then [ vlo ]\n else\n List.init n (fun i ->\n vlo +. (range *. (float (i + 1) /. float (n + 1))))\n in\n let contours = trace_contours ~rows ~cols data level_values in\n (* Map grid coords to data coords *)\n let xscale = (x1 -. x0) /. float (cols - 1) in\n let yscale = (y1 -. y0) /. float (rows - 1) in\n List.map\n (fun cp ->\n let paths =\n List.map\n (fun seg ->\n Array.map\n (fun (gc, gr) -> (x0 +. (gc *. xscale), y0 +. (gr *. yscale)))\n seg)\n cp.paths\n in\n { cp with paths })\n contours\n\n(* Compile a spec tree into a prepared tree *)\n\nlet compile_panel theme spec =\n let c = collect empty_collected spec in\n let c = { c with marks = List.rev c.marks } in\n let theme = Option.value ~default:theme c.theme_override in\n let marks =\n normalize_hist (normalize_imshow theme (auto_color theme c.marks))\n in\n let xscale = Option.value ~default:`Linear c.x.scale in\n let yscale = Option.value ~default:`Linear c.y.scale in\n let xlo, xhi, ylo, yhi = compute_data_bounds ~xscale ~yscale marks in\n let x = Axis.resolve ~data_lo:xlo ~data_hi:xhi c.x in\n let y = Axis.resolve ~data_lo:ylo ~data_hi:yhi c.y in\n Panel\n {\n marks;\n x;\n y;\n title = c.title;\n legend_loc = c.legend_loc;\n legend_ncol = c.legend_ncol;\n grid_visible = c.grid_visible;\n frame_visible = c.frame_visible;\n theme_override = c.theme_override;\n colorbar_range = color_by_range marks;\n size_by_range = size_by_range marks;\n }\n\nlet rec compile ~theme spec =\n match spec with\n | Spec.Grid { rows; gap } ->\n let rows = List.map (List.map (compile ~theme)) rows in\n Grid { rows; gap }\n | Spec.Decorated { inner = Spec.Grid _ as g; decorations } ->\n let gd = extract_grid_decorations decorations in\n let theme = Option.value ~default:theme gd.gd_theme_override in\n let all_marks = auto_color theme (collect_all_marks g) in\n let inner = compile ~theme g in\n Decorated_grid { decorations = gd; inner; all_marks }\n | spec -> compile_panel theme spec\n"
},
{
"path": "packages/hugin/lib/prepared.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Data-only compilation stage.\n\n {b Internal module.} Compiles a {!Spec.t} tree into a {!t} tree with all\n data-dependent work done: decoration collection, histogram binning,\n auto-coloring, data-bound computation, imshow rasterization, contour\n tracing, and guide-range detection.\n\n The result is independent of output dimensions and can be resolved\n repeatedly at different sizes by {!Resolve.resolve_prepared}. *)\n\n(** {1:bounds Data bounds} *)\n\nval nx_finite_range : Nx.float32_t -> float * float\n(** [nx_finite_range arr] is [(lo, hi)] of the finite values in [arr]. *)\n\n(** {1:marks Mark introspection} *)\n\nval mark_color : Spec.mark -> Color.t option\n(** [mark_color m] is the color of [m], if set. *)\n\n(** {1:contour Contour tracing} *)\n\ntype contour_paths = { level : float; paths : (float * float) array list }\n(** The type for traced contour paths at a single iso-level. Coordinates are in\n data space. *)\n\nval prepare_contour :\n x0:float ->\n x1:float ->\n y0:float ->\n y1:float ->\n data:Nx.float32_t ->\n levels:[ `Num of int | `Values of float array ] ->\n contour_paths list\n(** [prepare_contour ~x0 ~x1 ~y0 ~y1 ~data ~levels] traces contour paths through\n [data] and maps grid coordinates to the data-space rectangle \\[[x0];[x1]\\] x\n \\[[y0];[y1]\\]. *)\n\n(** {1:panel Prepared panel} *)\n\ntype panel = {\n marks : Spec.mark list;\n x : Axis.t;\n y : Axis.t;\n title : string option;\n legend_loc : Spec.legend_loc option;\n legend_ncol : int;\n grid_visible : bool option;\n frame_visible : bool option;\n theme_override : Theme.t option;\n colorbar_range : (float * float) option;\n size_by_range : (float * float) option;\n}\n(** The type for prepared panels. All data-only work is done: marks are\n auto-colored and histograms normalized to bars, data bounds are computed,\n and guide ranges are detected. *)\n\n(** {1:grid Grid decorations} *)\n\ntype grid_decorations = {\n gd_title : string option;\n gd_xlabel : string option;\n gd_ylabel : string option;\n gd_legend_loc : Spec.legend_loc option;\n gd_legend_ncol : int;\n gd_theme_override : Theme.t option;\n}\n(** The type for grid-level decorations extracted from a decorated grid spec. *)\n\n(** {1:tree Prepared tree} *)\n\ntype t =\n | Panel of panel\n | Grid of { rows : t list list; gap : float }\n | Decorated_grid of {\n decorations : grid_decorations;\n inner : t;\n all_marks : Spec.mark list;\n }\n (** The type for prepared spec trees. Mirrors {!Spec.t} structure with all\n data-only work pre-computed. *)\n\n(** {1:compile Compilation} *)\n\nval compile : theme:Theme.t -> Spec.t -> t\n(** [compile ~theme spec] is the prepared tree for [spec]. Collects decorations,\n normalizes histograms, auto-colors marks, computes data bounds, and detects\n colorbar/size-guide ranges. *)\n"
},
{
"path": "packages/hugin/lib/resolve.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Prepared.t + Theme.t -> Scene.t resolution\n\n Layout and pixel-coordinate work. All data-only processing (histogram\n binning, auto-coloring, bounds) is done in Prepared. *)\n\ntype text_measurer = font:Theme.font -> string -> float * float\n\n(* Region in device pixels *)\ntype region = { rx : float; ry : float; rw : float; rh : float }\n\n(* Scale-aware coord transform *)\n\nlet data_to_pixel_x sx region v =\n let u = sx.Scale.to_unit v in\n region.rx +. (u *. region.rw)\n\nlet data_to_pixel_y sy region v =\n let u = sy.Scale.to_unit v in\n region.ry +. region.rh -. (u *. region.rh)\n\n(* Dash pattern *)\n\nlet dash_of_style = function\n | `Solid -> []\n | `Dashed -> [ 6.; 4. ]\n | `Dotted -> [ 2.; 3. ]\n | `Dash_dot -> [ 6.; 3.; 2.; 3. ]\n\n(* Resolution helpers *)\n\nlet resolve_color ?default_alpha color alpha =\n let c = Option.value ~default:Color.black color in\n match (alpha, default_alpha) with\n | Some a, _ | None, Some a -> Color.with_alpha a c\n | None, None -> c\n\nlet resolve_line_width sf (theme : Theme.t) line_width =\n Option.value ~default:theme.line_width line_width *. sf\n\nlet resolve_dash sf line_style =\n let dash = match line_style with Some s -> dash_of_style s | None -> [] in\n List.map (fun d -> d *. sf) dash\n\n(* Emit mark primitives *)\n\nlet step_transform step points n =\n match step with\n | None -> points\n | Some `Post ->\n if n < 2 then points\n else begin\n let out = Array.make ((2 * n) - 1) (0., 0.) in\n let k = ref 0 in\n for i = 0 to n - 2 do\n let px, py = points.(i) in\n let px_next, _ = points.(i + 1) in\n out.(!k) <- (px, py);\n incr k;\n out.(!k) <- (px_next, py);\n incr k\n done;\n out.(!k) <- points.(n - 1);\n Array.sub out 0 (!k + 1)\n end\n | Some `Pre ->\n if n < 2 then points\n else begin\n let out = Array.make ((2 * n) - 1) (0., 0.) in\n let k = ref 0 in\n out.(!k) <- points.(0);\n incr k;\n for i = 1 to n - 1 do\n let _, py_prev = points.(i - 1) in\n let px, py = points.(i) in\n out.(!k) <- (px, py_prev);\n incr k;\n out.(!k) <- (px, py);\n incr k\n done;\n Array.sub out 0 !k\n end\n | Some `Mid ->\n if n < 2 then points\n else begin\n let out = Array.make ((3 * n) - 2) (0., 0.) in\n let k = ref 0 in\n for i = 0 to n - 2 do\n let px, py = points.(i) in\n let px_next, py_next = points.(i + 1) in\n let mx = (px +. px_next) /. 2. in\n out.(!k) <- (px, py);\n incr k;\n out.(!k) <- (mx, py);\n incr k;\n out.(!k) <- (mx, py_next);\n incr k\n done;\n out.(!k) <- points.(n - 1);\n Array.sub out 0 (!k + 1)\n end\n\nlet emit_line_mark sx sy plot_area theme ~x ~y ~color ~line_width ~line_style\n ~step ~marker ~alpha =\n let n = (Nx.shape x).(0) in\n let color = resolve_color color alpha in\n let sf = theme.Theme.scale_factor in\n let lw = resolve_line_width sf theme line_width in\n let scaled_dash = resolve_dash sf line_style in\n (* Split line into finite-value segments *)\n let segments = ref [] in\n let current = ref [] in\n let all_finite_points = ref [] in\n for i = 0 to n - 1 do\n let xv = Nx.item [ i ] x in\n let yv = Nx.item [ i ] y in\n if Float.is_finite xv && Float.is_finite yv then begin\n let px = data_to_pixel_x sx plot_area xv in\n let py = data_to_pixel_y sy plot_area yv in\n let pt = (px, py) in\n current := pt :: !current;\n all_finite_points := pt :: !all_finite_points\n end\n else\n match !current with\n | [] -> ()\n | _ ->\n segments := Array.of_list (List.rev !current) :: !segments;\n current := []\n done;\n (match !current with\n | [] -> ()\n | _ -> segments := Array.of_list (List.rev !current) :: !segments);\n let segments = List.rev !segments in\n let paths =\n List.map\n (fun points ->\n let n_pts = Array.length points in\n let points = step_transform step points n_pts in\n Scene.Path\n {\n points;\n close = false;\n fill = None;\n stroke = Some color;\n line_width = lw;\n dash = scaled_dash;\n })\n segments\n in\n match marker with\n | Some shape ->\n let finite_points = Array.of_list (List.rev !all_finite_points) in\n let ms = theme.marker_size *. sf in\n let markers =\n Scene.Markers\n {\n points = finite_points;\n shape;\n size = ms;\n sizes = None;\n fill = Some color;\n fills = None;\n stroke = None;\n }\n in\n paths @ [ markers ]\n | None -> paths\n\nlet emit_point_mark sx sy plot_area theme ~x ~y ~color ~color_by ~size ~size_by\n ~marker ~alpha =\n let n = (Nx.shape x).(0) in\n let color = resolve_color color alpha in\n let shape = Option.value ~default:Spec.Circle marker in\n let ms =\n (match size with Some s -> s | None -> theme.Theme.marker_size)\n *. theme.scale_factor\n in\n (* Collect only finite points *)\n let valid = Array.make n false in\n let num_valid = ref 0 in\n for i = 0 to n - 1 do\n let xv = Nx.item [ i ] x in\n let yv = Nx.item [ i ] y in\n if Float.is_finite xv && Float.is_finite yv then begin\n valid.(i) <- true;\n incr num_valid\n end\n done;\n let nv = !num_valid in\n let points = Array.make nv (0., 0.) in\n let vi = ref 0 in\n for i = 0 to n - 1 do\n if valid.(i) then begin\n let px = data_to_pixel_x sx plot_area (Nx.item [ i ] x) in\n let py = data_to_pixel_y sy plot_area (Nx.item [ i ] y) in\n points.(!vi) <- (px, py);\n incr vi\n end\n done;\n let sizes =\n match size_by with\n | Some sb ->\n let sb_lo, sb_hi = Prepared.nx_finite_range sb in\n let sb_range = if sb_hi = sb_lo then 1. else sb_hi -. sb_lo in\n let arr = Array.make nv ms in\n let vi = ref 0 in\n for i = 0 to n - 1 do\n if valid.(i) then begin\n let sv = (Nx.item [ i ] sb -. sb_lo) /. sb_range in\n arr.(!vi) <- (ms *. 0.5) +. (ms *. Float.sqrt sv);\n incr vi\n end\n done;\n Some arr\n | None -> None\n in\n let fills =\n match color_by with\n | Some cb ->\n let cb_lo, cb_hi = Prepared.nx_finite_range cb in\n let cb_range = if cb_hi = cb_lo then 1. else cb_hi -. cb_lo in\n let arr = Array.make nv Color.black in\n let vi = ref 0 in\n for i = 0 to n - 1 do\n if valid.(i) then begin\n let cv = (Nx.item [ i ] cb -. cb_lo) /. cb_range in\n let c = Cmap.eval theme.sequential cv in\n arr.(!vi) <-\n (match alpha with Some a -> Color.with_alpha a c | None -> c);\n incr vi\n end\n done;\n Some arr\n | None -> None\n in\n let fill = if fills <> None then None else Some color in\n let stroke = Some color in\n [ Scene.Markers { points; shape; size = ms; sizes; fill; fills; stroke } ]\n\nlet emit_bar_mark sx sy plot_area theme ~x ~height ~width ~bottom ~color ~alpha\n =\n let n = (Nx.shape x).(0) in\n let color = resolve_color color alpha in\n let w = Option.value ~default:0.8 width in\n let prims = ref [] in\n for i = 0 to n - 1 do\n let xi = Nx.item [ i ] x in\n let hi = Nx.item [ i ] height in\n if Float.is_finite xi && Float.is_finite hi then\n let x0 = data_to_pixel_x sx plot_area (xi -. (w /. 2.)) in\n let x1 = data_to_pixel_x sx plot_area (xi +. (w /. 2.)) in\n let y0 = data_to_pixel_y sy plot_area bottom in\n let y1 = data_to_pixel_y sy plot_area (bottom +. hi) in\n let lx = Float.min x0 x1 and rx = Float.max x0 x1 in\n let ty = Float.min y0 y1 and by = Float.max y0 y1 in\n prims :=\n Scene.Path\n {\n points = [| (lx, ty); (rx, ty); (rx, by); (lx, by) |];\n close = true;\n fill = Some color;\n stroke = None;\n line_width = 0.;\n dash = [];\n }\n :: !prims\n done;\n List.rev !prims\n\nlet emit_image_mark sx sy plot_area ~data ~extent =\n match extent with\n | Some (xmin, xmax, ymin, ymax) ->\n let px0 = data_to_pixel_x sx plot_area xmin in\n let px1 = data_to_pixel_x sx plot_area xmax in\n let py0 = data_to_pixel_y sy plot_area ymax in\n let py1 = data_to_pixel_y sy plot_area ymin in\n let x = Float.min px0 px1 in\n let y = Float.min py0 py1 in\n let w = Float.abs (px1 -. px0) in\n let h = Float.abs (py1 -. py0) in\n [ Scene.Image { x; y; w; h; data } ]\n | None ->\n [\n Scene.Image\n {\n x = plot_area.rx;\n y = plot_area.ry;\n w = plot_area.rw;\n h = plot_area.rh;\n data;\n };\n ]\n\nlet emit_text_mark sx sy plot_area theme ~x ~y ~content ~color ~font_size =\n let color = Option.value ~default:Color.black color in\n let size = Option.value ~default:theme.Theme.font_label.size font_size in\n let px = data_to_pixel_x sx plot_area x in\n let py = data_to_pixel_y sy plot_area y in\n [\n Scene.Text\n {\n x = px;\n y = py;\n content;\n font =\n {\n family = theme.font_label.family;\n size = size *. theme.scale_factor;\n weight = `Normal;\n };\n color;\n anchor = `Start;\n baseline = `Bottom;\n angle = 0.;\n };\n ]\n\nlet emit_hline_mark sy plot_area theme ~y:yv ~color ~line_width ~line_style\n ~alpha =\n let color = resolve_color color alpha in\n let sf = theme.Theme.scale_factor in\n let lw = resolve_line_width sf theme line_width in\n let dash = resolve_dash sf line_style in\n let py = data_to_pixel_y sy plot_area yv in\n [\n Scene.Path\n {\n points = [| (plot_area.rx, py); (plot_area.rx +. plot_area.rw, py) |];\n close = false;\n fill = None;\n stroke = Some color;\n line_width = lw;\n dash;\n };\n ]\n\nlet emit_vline_mark sx plot_area theme ~x:xv ~color ~line_width ~line_style\n ~alpha =\n let color = resolve_color color alpha in\n let sf = theme.Theme.scale_factor in\n let lw = resolve_line_width sf theme line_width in\n let dash = resolve_dash sf line_style in\n let px = data_to_pixel_x sx plot_area xv in\n [\n Scene.Path\n {\n points = [| (px, plot_area.ry); (px, plot_area.ry +. plot_area.rh) |];\n close = false;\n fill = None;\n stroke = Some color;\n line_width = lw;\n dash;\n };\n ]\n\nlet emit_abline_mark sx sy plot_area theme ~slope ~intercept ~color ~line_width\n ~line_style ~alpha =\n let color = resolve_color color alpha in\n let sf = theme.Theme.scale_factor in\n let lw = resolve_line_width sf theme line_width in\n let dash = resolve_dash sf line_style in\n let x0 = sx.Scale.lo and x1 = sx.Scale.hi in\n let y0v = (slope *. x0) +. intercept in\n let y1v = (slope *. x1) +. intercept in\n let px0 = data_to_pixel_x sx plot_area x0 in\n let py0 = data_to_pixel_y sy plot_area y0v in\n let px1 = data_to_pixel_x sx plot_area x1 in\n let py1 = data_to_pixel_y sy plot_area y1v in\n [\n Scene.Path\n {\n points = [| (px0, py0); (px1, py1) |];\n close = false;\n fill = None;\n stroke = Some color;\n line_width = lw;\n dash;\n };\n ]\n\nlet emit_fill_between_segment sx sy plot_area color indices x y1 y2 =\n let n_seg = List.length indices in\n if n_seg = 0 then []\n else\n let points = Array.make (2 * n_seg) (0., 0.) in\n let k = ref 0 in\n List.iter\n (fun i ->\n let xv = Nx.item [ i ] x in\n let yv = Nx.item [ i ] y1 in\n if Float.is_finite xv && Float.is_finite yv then begin\n points.(!k) <-\n (data_to_pixel_x sx plot_area xv, data_to_pixel_y sy plot_area yv);\n incr k\n end)\n indices;\n let forward_count = !k in\n List.iter\n (fun i ->\n let xv = Nx.item [ i ] x in\n let yv = Nx.item [ i ] y2 in\n if Float.is_finite xv && Float.is_finite yv then begin\n points.(!k) <-\n (data_to_pixel_x sx plot_area xv, data_to_pixel_y sy plot_area yv);\n incr k\n end)\n (List.rev indices);\n let total = !k in\n if total < 3 || forward_count = 0 then []\n else\n [\n Scene.Path\n {\n points = Array.sub points 0 total;\n close = true;\n fill = Some color;\n stroke = None;\n line_width = 0.;\n dash = [];\n };\n ]\n\nlet emit_fill_between_mark sx sy plot_area ~x ~y1 ~y2 ~where ~color ~alpha =\n let n = (Nx.shape x).(0) in\n let color = resolve_color ~default_alpha:0.3 color alpha in\n match where with\n | None ->\n let indices = List.init n Fun.id in\n emit_fill_between_segment sx sy plot_area color indices x y1 y2\n | Some mask ->\n (* Split into contiguous runs where mask > 0 *)\n let segments = ref [] in\n let current = ref [] in\n for i = 0 to n - 1 do\n if Nx.item [ i ] mask > 0. then current := i :: !current\n else\n match !current with\n | [] -> ()\n | seg ->\n segments := List.rev seg :: !segments;\n current := []\n done;\n (match !current with\n | [] -> ()\n | seg -> segments := List.rev seg :: !segments);\n List.concat_map\n (fun seg -> emit_fill_between_segment sx sy plot_area color seg x y1 y2)\n (List.rev !segments)\n\nlet emit_hspan_mark sy plot_area ~y0 ~y1 ~color ~alpha =\n let color = resolve_color ~default_alpha:0.2 color alpha in\n let py0 = data_to_pixel_y sy plot_area y0 in\n let py1 = data_to_pixel_y sy plot_area y1 in\n let top = Float.min py0 py1 and bot = Float.max py0 py1 in\n [\n Scene.Path\n {\n points =\n [|\n (plot_area.rx, top);\n (plot_area.rx +. plot_area.rw, top);\n (plot_area.rx +. plot_area.rw, bot);\n (plot_area.rx, bot);\n |];\n close = true;\n fill = Some color;\n stroke = None;\n line_width = 0.;\n dash = [];\n };\n ]\n\nlet emit_vspan_mark sx plot_area ~x0 ~x1 ~color ~alpha =\n let color = resolve_color ~default_alpha:0.2 color alpha in\n let px0 = data_to_pixel_x sx plot_area x0 in\n let px1 = data_to_pixel_x sx plot_area x1 in\n let left = Float.min px0 px1 and right = Float.max px0 px1 in\n [\n Scene.Path\n {\n points =\n [|\n (left, plot_area.ry);\n (right, plot_area.ry);\n (right, plot_area.ry +. plot_area.rh);\n (left, plot_area.ry +. plot_area.rh);\n |];\n close = true;\n fill = Some color;\n stroke = None;\n line_width = 0.;\n dash = [];\n };\n ]\n\nlet emit_errorbar_mark sx sy plot_area theme ~x ~y ~yerr ~xerr ~color\n ~line_width ~cap_size ~alpha =\n let n = (Nx.shape x).(0) in\n let color = resolve_color color alpha in\n let sf = theme.Theme.scale_factor in\n let lw = resolve_line_width sf theme line_width in\n let cap =\n (match cap_size with Some s -> s | None -> theme.marker_size *. 0.5) *. sf\n in\n let prims = ref [] in\n let make_path pts =\n Scene.Path\n {\n points = pts;\n close = false;\n fill = None;\n stroke = Some color;\n line_width = lw;\n dash = [];\n }\n in\n for i = 0 to n - 1 do\n let xv = Nx.item [ i ] x in\n let yv = Nx.item [ i ] y in\n if Float.is_finite xv && Float.is_finite yv then begin\n let px = data_to_pixel_x sx plot_area xv in\n let py = data_to_pixel_y sy plot_area yv in\n let y_lo, y_hi =\n match yerr with\n | `Symmetric e ->\n let ev = Nx.item [ i ] e in\n (yv -. ev, yv +. ev)\n | `Asymmetric (elo, ehi) ->\n (yv -. Nx.item [ i ] elo, yv +. Nx.item [ i ] ehi)\n in\n let py_lo = data_to_pixel_y sy plot_area y_lo in\n let py_hi = data_to_pixel_y sy plot_area y_hi in\n prims := make_path [| (px, py_lo); (px, py_hi) |] :: !prims;\n prims := make_path [| (px -. cap, py_hi); (px +. cap, py_hi) |] :: !prims;\n prims := make_path [| (px -. cap, py_lo); (px +. cap, py_lo) |] :: !prims;\n begin match xerr with\n | Some xerr_val ->\n let x_lo, x_hi =\n match xerr_val with\n | `Symmetric e ->\n let ev = Nx.item [ i ] e in\n (xv -. ev, xv +. ev)\n | `Asymmetric (elo, ehi) ->\n (xv -. Nx.item [ i ] elo, xv +. Nx.item [ i ] ehi)\n in\n let px_lo = data_to_pixel_x sx plot_area x_lo in\n let px_hi = data_to_pixel_x sx plot_area x_hi in\n prims := make_path [| (px_lo, py); (px_hi, py) |] :: !prims;\n prims :=\n make_path [| (px_lo, py -. cap); (px_lo, py +. cap) |] :: !prims;\n prims :=\n make_path [| (px_hi, py -. cap); (px_hi, py +. cap) |] :: !prims\n | None -> ()\n end\n end\n done;\n List.rev !prims\n\nlet emit_heatmap_mark sx sy plot_area theme ~data ~cmap ~annotate ~vmin ~vmax\n ~fmt =\n let shape = Nx.shape data in\n let rows = shape.(0) and cols = shape.(1) in\n let frows = float rows in\n let lo = ref Float.infinity and hi = ref Float.neg_infinity in\n for r = 0 to rows - 1 do\n for c = 0 to cols - 1 do\n let v = Nx.item [ r; c ] data in\n if Float.is_finite v then begin\n if v < !lo then lo := v;\n if v > !hi then hi := v\n end\n done\n done;\n let vlo = Option.value ~default:!lo vmin in\n let vhi = Option.value ~default:!hi vmax in\n let vrange = if vhi = vlo then 1. else vhi -. vlo in\n let cmap = Option.value ~default:theme.Theme.sequential cmap in\n let sf = theme.Theme.scale_factor in\n let prims = ref [] in\n for r = 0 to rows - 1 do\n for c = 0 to cols - 1 do\n let v = Nx.item [ r; c ] data in\n let t = Float.max 0. (Float.min 1. ((v -. vlo) /. vrange)) in\n let cell_color = Cmap.eval cmap t in\n let x0 = data_to_pixel_x sx plot_area (float c) in\n let x1 = data_to_pixel_x sx plot_area (float (c + 1)) in\n let y0 = data_to_pixel_y sy plot_area (frows -. float r) in\n let y1 = data_to_pixel_y sy plot_area (frows -. float (r + 1)) in\n let lx = Float.min x0 x1 and rx = Float.max x0 x1 in\n let ty = Float.min y0 y1 and by = Float.max y0 y1 in\n prims :=\n Scene.Path\n {\n points = [| (lx, ty); (rx, ty); (rx, by); (lx, by) |];\n close = true;\n fill = Some cell_color;\n stroke = None;\n line_width = 0.;\n dash = [];\n }\n :: !prims;\n if annotate then begin\n let text =\n match fmt with Some f -> f v | None -> Printf.sprintf \"%.2g\" v\n in\n let text_color =\n if Color.lightness cell_color > 0.65 then Color.black else Color.white\n in\n let cx = (lx +. rx) /. 2. in\n let cy = (ty +. by) /. 2. in\n let font_size =\n Float.max (8. *. sf)\n (Float.min\n (Float.abs (rx -. lx) *. 0.4)\n (Float.abs (by -. ty) *. 0.4))\n in\n prims :=\n Scene.Text\n {\n x = cx;\n y = cy;\n content = text;\n font =\n {\n family = theme.font_tick.family;\n size = font_size;\n weight = `Normal;\n };\n color = text_color;\n anchor = `Middle;\n baseline = `Middle;\n angle = 0.;\n }\n :: !prims\n end\n done\n done;\n List.rev !prims\n\nlet emit_contour_mark sx sy plot_area theme ~data ~x0 ~x1 ~y0 ~y1 ~levels\n ~filled ~cmap ~color ~line_width ~alpha =\n let sf = theme.Theme.scale_factor in\n let contours = Prepared.prepare_contour ~x0 ~x1 ~y0 ~y1 ~data ~levels in\n let n_levels = List.length contours in\n let prims = ref [] in\n List.iteri\n (fun i cp ->\n let t = if n_levels <= 1 then 0.5 else float i /. float (n_levels - 1) in\n let c =\n match color with\n | Some c -> c\n | None ->\n let cmap = Option.value ~default:theme.Theme.sequential cmap in\n Cmap.eval cmap t\n in\n let c = match alpha with Some a -> Color.with_alpha a c | None -> c in\n let lw = resolve_line_width sf theme line_width in\n List.iter\n (fun seg ->\n let points =\n Array.map\n (fun (dx, dy) ->\n ( data_to_pixel_x sx plot_area dx,\n data_to_pixel_y sy plot_area dy ))\n seg\n in\n if filled then\n prims :=\n Scene.Path\n {\n points;\n close = true;\n fill = Some c;\n stroke = None;\n line_width = 0.;\n dash = [];\n }\n :: !prims\n else\n prims :=\n Scene.Path\n {\n points;\n close = false;\n fill = None;\n stroke = Some c;\n line_width = lw;\n dash = [];\n }\n :: !prims)\n cp.Prepared.paths)\n contours;\n List.rev !prims\n\nlet emit_mark sx sy plot_area theme = function\n | Spec.Line\n { x; y; color; line_width; line_style; step; marker; label = _; alpha } ->\n emit_line_mark sx sy plot_area theme ~x ~y ~color ~line_width ~line_style\n ~step ~marker ~alpha\n | Spec.Point\n { x; y; color; color_by; size; size_by; marker; label = _; alpha } ->\n emit_point_mark sx sy plot_area theme ~x ~y ~color ~color_by ~size\n ~size_by ~marker ~alpha\n | Spec.Bar { x; height; width; bottom; color; label = _; alpha } ->\n emit_bar_mark sx sy plot_area theme ~x ~height ~width ~bottom ~color\n ~alpha\n | Spec.Hist _ ->\n failwith\n \"resolve: Spec.Hist reached emit_mark; should have been normalized to \\\n Bar by Prepared.compile\"\n | Spec.Image { data; extent } -> emit_image_mark sx sy plot_area ~data ~extent\n | Spec.Text_mark { x; y; content; color; font_size } ->\n emit_text_mark sx sy plot_area theme ~x ~y ~content ~color ~font_size\n | Spec.Hline { y; color; line_width; line_style; label = _; alpha } ->\n emit_hline_mark sy plot_area theme ~y ~color ~line_width ~line_style\n ~alpha\n | Spec.Vline { x; color; line_width; line_style; label = _; alpha } ->\n emit_vline_mark sx plot_area theme ~x ~color ~line_width ~line_style\n ~alpha\n | Spec.Abline\n { slope; intercept; color; line_width; line_style; label = _; alpha } ->\n emit_abline_mark sx sy plot_area theme ~slope ~intercept ~color\n ~line_width ~line_style ~alpha\n | Spec.Fill_between { x; y1; y2; where; color; alpha; label = _ } ->\n emit_fill_between_mark sx sy plot_area ~x ~y1 ~y2 ~where ~color ~alpha\n | Spec.Hspan { y0; y1; color; alpha; label = _ } ->\n emit_hspan_mark sy plot_area ~y0 ~y1 ~color ~alpha\n | Spec.Vspan { x0; x1; color; alpha; label = _ } ->\n emit_vspan_mark sx plot_area ~x0 ~x1 ~color ~alpha\n | Spec.Errorbar\n { x; y; yerr; xerr; color; line_width; cap_size; label = _; alpha } ->\n emit_errorbar_mark sx sy plot_area theme ~x ~y ~yerr ~xerr ~color\n ~line_width ~cap_size ~alpha\n | Spec.Heatmap { data; cmap; annotate; vmin; vmax; fmt } ->\n emit_heatmap_mark sx sy plot_area theme ~data ~cmap ~annotate ~vmin ~vmax\n ~fmt\n | Spec.Imshow _ ->\n failwith\n \"resolve: Spec.Imshow reached emit_mark; should have been normalized \\\n to Image by Prepared.compile\"\n | Spec.Contour\n {\n data;\n x0;\n x1;\n y0;\n y1;\n levels;\n filled;\n cmap;\n color;\n line_width;\n label = _;\n alpha;\n } ->\n emit_contour_mark sx sy plot_area theme ~data ~x0 ~x1 ~y0 ~y1 ~levels\n ~filled ~cmap ~color ~line_width ~alpha\n\n(* Axis primitives *)\n\nlet scaled_font (theme : Theme.t) (f : Theme.font) =\n { f with size = f.size *. theme.scale_factor }\n\nlet emit_axes ~text_measurer sx sy plot_area (theme : Theme.t) ~xticks ~yticks\n (pp : Prepared.panel) =\n let sf = theme.scale_factor in\n let prims = ref [] in\n let axis_color = theme.axis.color in\n let lw = theme.axis.width *. sf in\n let tl = theme.tick_length *. sf in\n List.iter\n (fun (v, label) ->\n let px = data_to_pixel_x sx plot_area v in\n let by = plot_area.ry +. plot_area.rh in\n prims :=\n Scene.Path\n {\n points = [| (px, by); (px, by +. tl) |];\n close = false;\n fill = None;\n stroke = Some axis_color;\n line_width = lw;\n dash = [];\n }\n :: !prims;\n let font = scaled_font theme theme.font_tick in\n prims :=\n Scene.Text\n {\n x = px;\n y = by +. tl +. (8. *. sf);\n content = label;\n font;\n color = axis_color;\n anchor = `Middle;\n baseline = `Top;\n angle = 0.;\n }\n :: !prims)\n xticks;\n\n (* Y ticks *)\n List.iter\n (fun (v, label) ->\n let py = data_to_pixel_y sy plot_area v in\n let lx = plot_area.rx in\n prims :=\n Scene.Path\n {\n points = [| (lx -. tl, py); (lx, py) |];\n close = false;\n fill = None;\n stroke = Some axis_color;\n line_width = lw;\n dash = [];\n }\n :: !prims;\n let font = scaled_font theme theme.font_tick in\n prims :=\n Scene.Text\n {\n x = lx -. tl -. (8. *. sf);\n y = py;\n content = label;\n font;\n color = axis_color;\n anchor = `End;\n baseline = `Middle;\n angle = 0.;\n }\n :: !prims)\n yticks;\n\n (* Grid *)\n let show_grid = Option.value ~default:(theme.grid <> None) pp.grid_visible in\n begin match theme.grid with\n | Some grid_line when show_grid ->\n List.iter\n (fun (v, _) ->\n let px = data_to_pixel_x sx plot_area v in\n prims :=\n Scene.Path\n {\n points =\n [| (px, plot_area.ry); (px, plot_area.ry +. plot_area.rh) |];\n close = false;\n fill = None;\n stroke = Some grid_line.color;\n line_width = grid_line.width *. sf;\n dash = grid_line.dash;\n }\n :: !prims)\n xticks;\n List.iter\n (fun (v, _) ->\n let py = data_to_pixel_y sy plot_area v in\n prims :=\n Scene.Path\n {\n points =\n [| (plot_area.rx, py); (plot_area.rx +. plot_area.rw, py) |];\n close = false;\n fill = None;\n stroke = Some grid_line.color;\n line_width = grid_line.width *. sf;\n dash = grid_line.dash;\n }\n :: !prims)\n yticks\n | _ -> ()\n end;\n\n (* Axis border *)\n let show_frame = Option.value ~default:true pp.frame_visible in\n if show_frame then begin\n let lx = plot_area.rx and ty = plot_area.ry in\n let rx = lx +. plot_area.rw and by = ty +. plot_area.rh in\n prims :=\n Scene.Path\n {\n points = [| (lx, ty); (rx, ty); (rx, by); (lx, by) |];\n close = true;\n fill = None;\n stroke = Some axis_color;\n line_width = lw;\n dash = [];\n }\n :: !prims\n end;\n\n (* Title *)\n begin match pp.title with\n | Some s ->\n let font = scaled_font theme theme.font_title in\n let cx = plot_area.rx +. (plot_area.rw /. 2.) in\n prims :=\n Scene.Text\n {\n x = cx;\n y = plot_area.ry -. (theme.title_gap *. sf);\n content = s;\n font;\n color = axis_color;\n anchor = `Middle;\n baseline = `Bottom;\n angle = 0.;\n }\n :: !prims\n | None -> ()\n end;\n\n (* X label *)\n begin match pp.x.label with\n | Some s ->\n let font = scaled_font theme theme.font_label in\n let cx = plot_area.rx +. (plot_area.rw /. 2.) in\n let _, tick_h =\n text_measurer ~font:(scaled_font theme theme.font_tick) \"0\"\n in\n let y =\n plot_area.ry +. plot_area.rh +. tl +. tick_h +. (theme.label_gap *. sf)\n in\n prims :=\n Scene.Text\n {\n x = cx;\n y;\n content = s;\n font;\n color = axis_color;\n anchor = `Middle;\n baseline = `Top;\n angle = 0.;\n }\n :: !prims\n | None -> ()\n end;\n\n (* Y label *)\n begin match pp.y.label with\n | Some s ->\n let font = scaled_font theme theme.font_label in\n let tick_font = scaled_font theme theme.font_tick in\n let max_ytick_w =\n List.fold_left\n (fun acc (_, label) ->\n let w, _ = text_measurer ~font:tick_font label in\n Float.max acc w)\n 0. yticks\n in\n let _, label_h = text_measurer ~font s in\n let x =\n plot_area.rx -. tl -. max_ytick_w -. (8. *. sf)\n -. (theme.label_gap *. sf) -. (label_h /. 2.)\n in\n let y = plot_area.ry +. (plot_area.rh /. 2.) in\n prims :=\n Scene.Text\n {\n x;\n y;\n content = s;\n font;\n color = axis_color;\n anchor = `Middle;\n baseline = `Middle;\n angle = Float.pi /. 2.;\n }\n :: !prims\n | None -> ()\n end;\n\n List.rev !prims\n\n(* Legend *)\n\ntype legend_kind =\n | Legend_line of Spec.line_style option * Spec.marker option\n | Legend_point of Spec.marker\n | Legend_bar\n | Legend_ref_line of Spec.line_style option\n\nlet mark_label = function\n | Spec.Line { label; _ }\n | Spec.Point { label; _ }\n | Spec.Bar { label; _ }\n | Spec.Hist { label; _ }\n | Spec.Hline { label; _ }\n | Spec.Vline { label; _ }\n | Spec.Abline { label; _ }\n | Spec.Fill_between { label; _ }\n | Spec.Hspan { label; _ }\n | Spec.Vspan { label; _ }\n | Spec.Errorbar { label; _ }\n | Spec.Contour { label; _ } ->\n label\n | Spec.Image _ | Spec.Text_mark _ | Spec.Heatmap _ | Spec.Imshow _ -> None\n\nlet mark_legend_kind = function\n | Spec.Line { line_style; marker; _ } -> Legend_line (line_style, marker)\n | Spec.Point { marker; _ } ->\n Legend_point (Option.value ~default:Spec.Circle marker)\n | Spec.Bar _ | Spec.Hist _ | Spec.Fill_between _ | Spec.Hspan _ | Spec.Vspan _\n ->\n Legend_bar\n | Spec.Hline { line_style; _ }\n | Spec.Vline { line_style; _ }\n | Spec.Abline { line_style; _ } ->\n Legend_ref_line line_style\n | Spec.Errorbar _ -> Legend_ref_line None\n | Spec.Contour { filled = true; _ } -> Legend_bar\n | Spec.Contour _ -> Legend_ref_line None\n | _ -> Legend_bar\n\nlet emit_legend ~text_measurer ~loc ~ncol plot_area theme marks =\n let sf = theme.Theme.scale_factor in\n let entries =\n List.filter_map\n (fun m ->\n match mark_label m with\n | Some label ->\n let color =\n match Prepared.mark_color m with\n | Some c -> c\n | None -> Color.black\n in\n Some (label, color, mark_legend_kind m)\n | None -> None)\n marks\n in\n if entries = [] then []\n else begin\n let font = scaled_font theme theme.font_tick in\n let swatch_size = 10. *. sf in\n let gap = 4. *. sf in\n let line_h = Float.max swatch_size (font.size *. 1.2) in\n let margin = 8. *. sf in\n let ncol = max 1 ncol in\n let n_entries = List.length entries in\n let nrows = (n_entries + ncol - 1) / ncol in\n (* Compute per-column max label width *)\n let col_widths = Array.make ncol 0. in\n List.iteri\n (fun i (label, _, _) ->\n let col = i mod ncol in\n let w, _ = text_measurer ~font label in\n col_widths.(col) <- Float.max col_widths.(col) w)\n entries;\n let col_gap = 12. *. sf in\n let col_w i = swatch_size +. gap +. col_widths.(i) in\n let legend_w =\n let total = ref 0. in\n for i = 0 to ncol - 1 do\n total := !total +. col_w i\n done;\n !total +. (col_gap *. float (ncol - 1))\n in\n let legend_h = (float nrows *. (line_h +. gap)) -. gap in\n let loc = Option.value ~default:Spec.Upper_right loc in\n (* x0 is the right edge of the legend area *)\n let x0, y0 =\n match loc with\n | Spec.Upper_right ->\n (plot_area.rx +. plot_area.rw -. margin, plot_area.ry +. margin)\n | Spec.Upper_left ->\n (plot_area.rx +. margin +. legend_w, plot_area.ry +. margin)\n | Spec.Lower_right ->\n ( plot_area.rx +. plot_area.rw -. margin,\n plot_area.ry +. plot_area.rh -. margin -. legend_h )\n | Spec.Lower_left ->\n ( plot_area.rx +. margin +. legend_w,\n plot_area.ry +. plot_area.rh -. margin -. legend_h )\n | Spec.Center ->\n ( plot_area.rx +. ((plot_area.rw +. legend_w) /. 2.),\n plot_area.ry +. ((plot_area.rh -. legend_h) /. 2.) )\n | Spec.Right ->\n ( plot_area.rx +. plot_area.rw -. margin,\n plot_area.ry +. ((plot_area.rh -. legend_h) /. 2.) )\n | Spec.Upper_center ->\n ( plot_area.rx +. ((plot_area.rw +. legend_w) /. 2.),\n plot_area.ry +. margin )\n | Spec.Lower_center ->\n ( plot_area.rx +. ((plot_area.rw +. legend_w) /. 2.),\n plot_area.ry +. plot_area.rh -. margin -. legend_h )\n in\n (* Background box *)\n let inner_pad = 6. *. sf in\n let bg_x = x0 -. legend_w -. inner_pad in\n let bg_y = y0 -. inner_pad in\n let bg_w = legend_w +. (2. *. inner_pad) in\n let bg_h = legend_h +. (2. *. inner_pad) in\n let bg =\n Scene.Path\n {\n points =\n [|\n (bg_x, bg_y);\n (bg_x +. bg_w, bg_y);\n (bg_x +. bg_w, bg_y +. bg_h);\n (bg_x, bg_y +. bg_h);\n |];\n close = true;\n fill = Some (Color.with_alpha 0.85 theme.background);\n stroke = Some (Color.with_alpha 0.3 theme.axis.color);\n line_width = 1. *. sf;\n dash = [];\n }\n in\n (* Compute column x-offsets (from right edge of legend) *)\n let col_offsets = Array.make ncol 0. in\n let acc = ref 0. in\n for c = ncol - 1 downto 0 do\n col_offsets.(c) <- !acc;\n acc := !acc +. col_w c +. col_gap\n done;\n let prims = ref [ bg ] in\n List.iteri\n (fun i (label, color, kind) ->\n let row = i / ncol in\n let col = i mod ncol in\n let y = y0 +. (float row *. (line_h +. gap)) in\n let y_mid = y +. (swatch_size /. 2.) in\n let cx0 = x0 -. col_offsets.(col) in\n begin match kind with\n | Legend_line (line_style, marker) ->\n prims :=\n Scene.Path\n {\n points = [| (cx0 -. swatch_size, y_mid); (cx0, y_mid) |];\n close = false;\n fill = None;\n stroke = Some color;\n line_width = theme.line_width *. sf;\n dash = resolve_dash sf line_style;\n }\n :: !prims;\n begin match marker with\n | Some shape ->\n let ms = 6. *. sf in\n prims :=\n Scene.Markers\n {\n points = [| (cx0 -. (swatch_size /. 2.), y_mid) |];\n shape;\n size = ms;\n sizes = None;\n fill = Some color;\n fills = None;\n stroke = None;\n }\n :: !prims\n | None -> ()\n end\n | Legend_point marker ->\n let ms = 8. *. sf in\n prims :=\n Scene.Markers\n {\n points = [| (cx0 -. (swatch_size /. 2.), y_mid) |];\n shape = marker;\n size = ms;\n sizes = None;\n fill = Some color;\n fills = None;\n stroke = None;\n }\n :: !prims\n | Legend_bar ->\n prims :=\n Scene.Path\n {\n points =\n [|\n (cx0 -. swatch_size, y);\n (cx0, y);\n (cx0, y +. swatch_size);\n (cx0 -. swatch_size, y +. swatch_size);\n |];\n close = true;\n fill = Some color;\n stroke = None;\n line_width = 0.;\n dash = [];\n }\n :: !prims\n | Legend_ref_line line_style ->\n prims :=\n Scene.Path\n {\n points = [| (cx0 -. swatch_size, y_mid); (cx0, y_mid) |];\n close = false;\n fill = None;\n stroke = Some color;\n line_width = theme.line_width *. sf;\n dash = resolve_dash sf line_style;\n }\n :: !prims\n end;\n prims :=\n Scene.Text\n {\n x = cx0 -. swatch_size -. gap;\n y = y_mid;\n content = label;\n font;\n color = theme.axis.color;\n anchor = `End;\n baseline = `Middle;\n angle = 0.;\n }\n :: !prims)\n entries;\n List.rev !prims\n end\n\n(* Colorbar for color_by *)\n\nlet emit_colorbar plot_area (theme : Theme.t) ~height_frac (lo, hi) =\n let sf = theme.scale_factor in\n let font = scaled_font theme theme.font_tick in\n let bar_w = 16. *. sf in\n let bar_gap = 12. *. sf in\n let bar_x = plot_area.rx +. plot_area.rw +. bar_gap in\n let bar_y = plot_area.ry in\n let bar_h = plot_area.rh *. height_frac in\n (* Vertical gradient: series of thin horizontal strips *)\n let n_strips = 64 in\n let strip_h = bar_h /. float n_strips in\n let strips =\n List.init n_strips (fun i ->\n let t = 1. -. (float i /. float (n_strips - 1)) in\n let c = Cmap.eval theme.sequential t in\n let sy = bar_y +. (float i *. strip_h) in\n Scene.Path\n {\n points =\n [|\n (bar_x, sy);\n (bar_x +. bar_w, sy);\n (bar_x +. bar_w, sy +. strip_h +. 1.);\n (bar_x, sy +. strip_h +. 1.);\n |];\n close = true;\n fill = Some c;\n stroke = None;\n line_width = 0.;\n dash = [];\n })\n in\n (* Border around colorbar *)\n let border =\n Scene.Path\n {\n points =\n [|\n (bar_x, bar_y);\n (bar_x +. bar_w, bar_y);\n (bar_x +. bar_w, bar_y +. bar_h);\n (bar_x, bar_y +. bar_h);\n |];\n close = true;\n fill = None;\n stroke = Some theme.axis.color;\n line_width = theme.axis.width *. sf;\n dash = [];\n }\n in\n (* Tick labels along the right edge *)\n let ticks = Ticks.generate `Linear ~lo ~hi () in\n let range = hi -. lo in\n let range = if range = 0. then 1. else range in\n let label_x = bar_x +. bar_w +. (6. *. sf) in\n let tick_prims =\n List.filter_map\n (fun (v, label) ->\n let t = (v -. lo) /. range in\n if t < -0.01 || t > 1.01 then None\n else\n let py = bar_y +. bar_h -. (t *. bar_h) in\n Some\n (Scene.Text\n {\n x = label_x;\n y = py;\n content = label;\n font;\n color = theme.axis.color;\n anchor = `Start;\n baseline = `Middle;\n angle = 0.;\n }))\n ticks\n in\n strips @ [ border ] @ tick_prims\n\n(* Size guide for size_by *)\n\nlet emit_size_guide plot_area (theme : Theme.t) ~y_offset (lo, hi) =\n let sf = theme.scale_factor in\n let font = scaled_font theme theme.font_tick in\n let guide_gap = 12. *. sf in\n let guide_x = plot_area.rx +. plot_area.rw +. guide_gap in\n let max_r = 12. *. sf in\n (* Three representative sizes: max, mid, min *)\n let values = [| hi; (lo +. hi) /. 2.; lo |] in\n let range = hi -. lo in\n let range = if range = 0. then 1. else range in\n let prims = ref [] in\n let cy = ref (plot_area.ry +. y_offset +. max_r +. (4. *. sf)) in\n Array.iter\n (fun v ->\n let t = (v -. lo) /. range in\n let size = ((max_r *. 0.3) +. (max_r *. 0.7 *. Float.sqrt t)) *. 2. in\n let cx = guide_x +. max_r in\n prims :=\n Scene.Markers\n {\n points = [| (cx, !cy) |];\n shape = Spec.Circle;\n size;\n sizes = None;\n fill = Some (Color.with_alpha 0.2 theme.axis.color);\n fills = None;\n stroke = Some theme.axis.color;\n }\n :: !prims;\n let label = Printf.sprintf \"%.4g\" v in\n let label_x = cx +. max_r +. (6. *. sf) in\n prims :=\n Scene.Text\n {\n x = label_x;\n y = !cy;\n content = label;\n font;\n color = theme.axis.color;\n anchor = `Start;\n baseline = `Middle;\n angle = 0.;\n }\n :: !prims;\n cy := !cy +. (max_r *. 2.) +. (8. *. sf))\n values;\n List.rev !prims\n\n(* Compute layout padding *)\n\nlet compute_layout ~text_measurer (theme : Theme.t) (pp : Prepared.panel) xticks\n yticks =\n let sf = theme.scale_factor in\n let tick_font = scaled_font theme theme.font_tick in\n let label_font = scaled_font theme theme.font_label in\n let title_font = scaled_font theme theme.font_title in\n let tl = theme.tick_length *. sf in\n\n (* Left padding: y-tick labels + gap + optional ylabel *)\n let left =\n let base = theme.padding *. sf in\n match yticks with\n | [] -> base\n | _ ->\n let max_ytick_w =\n List.fold_left\n (fun acc (_, label) ->\n let w, _ = text_measurer ~font:tick_font label in\n Float.max acc w)\n 0. yticks\n in\n base +. max_ytick_w +. tl +. (8. *. sf)\n in\n let left =\n match pp.y.label with\n | Some s ->\n let _, h = text_measurer ~font:label_font s in\n left +. h +. (theme.label_gap *. sf)\n | None -> left\n in\n\n (* Bottom padding: x-tick labels + gap + optional xlabel *)\n let bottom =\n let base = theme.padding *. sf in\n match xticks with\n | [] -> base\n | _ ->\n let _, tick_h = text_measurer ~font:tick_font \"0\" in\n base +. tick_h +. tl +. (8. *. sf)\n in\n let bottom =\n match pp.x.label with\n | Some s ->\n let _, h = text_measurer ~font:label_font s in\n bottom +. h +. (theme.label_gap *. sf)\n | None -> bottom\n in\n\n (* Top padding: title *)\n let top = theme.padding *. sf in\n let top =\n match pp.title with\n | Some s ->\n let _, h = text_measurer ~font:title_font s in\n top +. h +. (theme.title_gap *. sf)\n | None -> top\n in\n\n (* Right padding — extra space for colorbar / size guide *)\n let right =\n let base = theme.padding *. sf in\n let colorbar_w =\n match pp.colorbar_range with\n | Some (lo, hi) ->\n let bar_w = 16. *. sf in\n let bar_gap = 12. *. sf in\n let ticks = Ticks.generate `Linear ~lo ~hi () in\n let max_label_w =\n List.fold_left\n (fun acc (_, label) ->\n let w, _ = text_measurer ~font:tick_font label in\n Float.max acc w)\n 0. ticks\n in\n bar_gap +. bar_w +. (6. *. sf) +. max_label_w +. (4. *. sf)\n | None -> 0.\n in\n let size_guide_w =\n match pp.size_by_range with\n | Some (lo, hi) ->\n let guide_gap = 12. *. sf in\n let max_r = 12. *. sf in\n let mid = (lo +. hi) /. 2. in\n let max_label_w =\n List.fold_left\n (fun acc v ->\n let w, _ =\n text_measurer ~font:tick_font (Printf.sprintf \"%.4g\" v)\n in\n Float.max acc w)\n 0. [ lo; mid; hi ]\n in\n guide_gap +. (max_r *. 2.) +. (6. *. sf) +. max_label_w +. (4. *. sf)\n | None -> 0.\n in\n base +. Float.max colorbar_w size_guide_w\n in\n\n (left, top, right, bottom)\n\n(* Resolve a single prepared panel *)\n\nlet resolve_panel ~text_measurer theme region (pp : Prepared.panel) =\n let theme = Option.value ~default:theme pp.theme_override in\n\n let sx, xticks = Axis.make_scale_and_ticks pp.x in\n let sy, yticks = Axis.make_scale_and_ticks pp.y in\n\n let left, top, right, bottom =\n compute_layout ~text_measurer theme pp xticks yticks\n in\n\n let plot_area =\n {\n rx = region.rx +. left;\n ry = region.ry +. top;\n rw = Float.max 1. (region.rw -. left -. right);\n rh = Float.max 1. (region.rh -. top -. bottom);\n }\n in\n\n (* Background *)\n let bg =\n Scene.Path\n {\n points =\n [|\n (region.rx, region.ry);\n (region.rx +. region.rw, region.ry);\n (region.rx +. region.rw, region.ry +. region.rh);\n (region.rx, region.ry +. region.rh);\n |];\n close = true;\n fill = Some theme.background;\n stroke = None;\n line_width = 0.;\n dash = [];\n }\n in\n\n (* Axes decorations *)\n let axes_prims =\n emit_axes ~text_measurer sx sy plot_area theme ~xticks ~yticks pp\n in\n\n (* Data marks inside clip region *)\n let data_prims = List.concat_map (emit_mark sx sy plot_area theme) pp.marks in\n let clipped_data =\n Scene.Clip\n {\n x = plot_area.rx;\n y = plot_area.ry;\n w = plot_area.rw;\n h = plot_area.rh;\n children = data_prims;\n }\n in\n\n (* Legend *)\n let legend_prims =\n emit_legend ~text_measurer ~loc:pp.legend_loc ~ncol:pp.legend_ncol plot_area\n theme pp.marks\n in\n\n (* Colorbar for color_by *)\n let has_both = pp.colorbar_range <> None && pp.size_by_range <> None in\n let colorbar_prims =\n match pp.colorbar_range with\n | Some range ->\n let height_frac = if has_both then 0.55 else 1. in\n emit_colorbar plot_area theme ~height_frac range\n | None -> []\n in\n\n (* Size guide for size_by *)\n let size_guide_prims =\n match pp.size_by_range with\n | Some range ->\n let y_offset = if has_both then plot_area.rh *. 0.6 else 0. in\n emit_size_guide plot_area theme ~y_offset range\n | None -> []\n in\n\n [ bg; clipped_data ] @ axes_prims @ legend_prims @ colorbar_prims\n @ size_guide_prims\n\n(* Resolve a prepared grid layout *)\n\nlet resolve_grid ~resolve_prepared ~text_measurer theme region rows gap =\n let nrows = List.length rows in\n let ncols =\n List.fold_left (fun acc row -> max acc (List.length row)) 0 rows\n in\n if nrows = 0 || ncols = 0 then []\n else begin\n let cell_w = (region.rw -. (gap *. float (ncols - 1))) /. float ncols in\n let cell_h = (region.rh -. (gap *. float (nrows - 1))) /. float nrows in\n let prims = ref [] in\n List.iteri\n (fun ri row ->\n List.iteri\n (fun ci prepared ->\n let cell_region =\n {\n rx = region.rx +. (float ci *. (cell_w +. gap));\n ry = region.ry +. (float ri *. (cell_h +. gap));\n rw = cell_w;\n rh = cell_h;\n }\n in\n let p =\n resolve_prepared ~text_measurer theme cell_region prepared\n in\n prims := List.rev_append p !prims)\n row)\n rows;\n List.rev !prims\n end\n\n(* Grid-level decorations *)\n\nlet emit_grid_decorations ~text_measurer theme region\n (gd : Prepared.grid_decorations) all_marks =\n let sf = theme.Theme.scale_factor in\n let color = theme.axis.color in\n let prims = ref [] in\n let r = ref region in\n\n (* Title: above grid *)\n begin match gd.gd_title with\n | Some s ->\n let font = scaled_font theme theme.font_title in\n let _, title_h = text_measurer ~font s in\n let title_gap = theme.title_gap *. sf in\n prims :=\n Scene.Text\n {\n x = !r.rx +. (!r.rw /. 2.);\n y = !r.ry +. title_h;\n content = s;\n font;\n color;\n anchor = `Middle;\n baseline = `Bottom;\n angle = 0.;\n }\n :: !prims;\n let used = title_h +. title_gap in\n r := { !r with ry = !r.ry +. used; rh = !r.rh -. used }\n | None -> ()\n end;\n\n (* Xlabel: below grid *)\n begin match gd.gd_xlabel with\n | Some s ->\n let font = scaled_font theme theme.font_label in\n let _, label_h = text_measurer ~font s in\n let label_gap = theme.label_gap *. sf in\n let used = label_h +. label_gap in\n prims :=\n Scene.Text\n {\n x = !r.rx +. (!r.rw /. 2.);\n y = !r.ry +. !r.rh -. label_gap;\n content = s;\n font;\n color;\n anchor = `Middle;\n baseline = `Bottom;\n angle = 0.;\n }\n :: !prims;\n r := { !r with rh = !r.rh -. used }\n | None -> ()\n end;\n\n (* Ylabel: left of grid, rotated *)\n begin match gd.gd_ylabel with\n | Some s ->\n let font = scaled_font theme theme.font_label in\n let _, label_h = text_measurer ~font s in\n let label_gap = theme.label_gap *. sf in\n let used = label_h +. label_gap in\n prims :=\n Scene.Text\n {\n x = !r.rx +. (label_h /. 2.);\n y = !r.ry +. (!r.rh /. 2.);\n content = s;\n font;\n color;\n anchor = `Middle;\n baseline = `Middle;\n angle = Float.pi /. 2.;\n }\n :: !prims;\n r := { !r with rx = !r.rx +. used; rw = !r.rw -. used }\n | None -> ()\n end;\n\n (* Shared legend *)\n let legend_prims =\n match gd.gd_legend_loc with\n | Some loc ->\n emit_legend ~text_measurer ~loc:(Some loc)\n ~ncol:gd.Prepared.gd_legend_ncol !r theme all_marks\n | None -> []\n in\n\n (List.rev !prims, legend_prims, !r)\n\n(* Top-level resolve from Prepared.t *)\n\nlet rec resolve_tree ~text_measurer theme region = function\n | Prepared.Panel pp -> resolve_panel ~text_measurer theme region pp\n | Prepared.Grid { rows; gap } ->\n let gap_px = gap *. Float.min region.rw region.rh in\n resolve_grid ~resolve_prepared:resolve_tree ~text_measurer theme region\n rows gap_px\n | Prepared.Decorated_grid { decorations; inner; all_marks } ->\n let theme = Option.value ~default:theme decorations.gd_theme_override in\n let dec_prims, legend_prims, grid_region =\n emit_grid_decorations ~text_measurer theme region decorations all_marks\n in\n dec_prims\n @ resolve_tree ~text_measurer theme grid_region inner\n @ legend_prims\n\nlet resolve_prepared ~text_measurer ~theme ~width ~height prepared =\n let region = { rx = 0.; ry = 0.; rw = width; rh = height } in\n let primitives = resolve_tree ~text_measurer theme region prepared in\n { Scene.width; height; primitives }\n\n(* Convenience: compile + resolve in one step *)\n\nlet resolve ~text_measurer ~theme ~width ~height spec =\n let prepared = Prepared.compile ~theme spec in\n resolve_prepared ~text_measurer ~theme ~width ~height prepared\n"
},
{
"path": "packages/hugin/lib/resolve.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Specification to scene resolution.\n\n {b Internal module.} Walks a {!Spec.t} tree, computes data bounds and\n layout, and emits a {!Scene.t} with all coordinates in device pixels. *)\n\ntype text_measurer = font:Theme.font -> string -> float * float\n(** The type for text measurers. Returns [(width, height)] for a string rendered\n in the given font. *)\n\nval resolve_prepared :\n text_measurer:text_measurer ->\n theme:Theme.t ->\n width:float ->\n height:float ->\n Prepared.t ->\n Scene.t\n(** [resolve_prepared ~text_measurer ~theme ~width ~height prepared] is the\n resolved scene for [prepared] at the given dimensions. Layout-only work\n (pixel coordinates, text measurement) is done here; data work is already\n done in {!Prepared.compile}. *)\n\nval resolve :\n text_measurer:text_measurer ->\n theme:Theme.t ->\n width:float ->\n height:float ->\n Spec.t ->\n Scene.t\n(** [resolve ~text_measurer ~theme ~width ~height spec] is the resolved scene\n for [spec] at the given dimensions. Convenience wrapper that calls\n {!Prepared.compile} then {!resolve_prepared}. *)\n"
},
{
"path": "packages/hugin/lib/scale.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Data-to-unit mapping functions *)\n\ntype t = {\n to_unit : float -> float;\n from_unit : float -> float;\n lo : float;\n hi : float;\n}\n\nlet maybe_invert invert to_unit from_unit =\n if invert then ((fun v -> 1. -. to_unit v), fun u -> from_unit (1. -. u))\n else (to_unit, from_unit)\n\nlet linear ?(invert = false) ~lo ~hi () =\n let range = hi -. lo in\n let range = if range = 0. then 1. else range in\n let to_unit, from_unit =\n maybe_invert invert\n (fun v -> (v -. lo) /. range)\n (fun u -> lo +. (u *. range))\n in\n { to_unit; from_unit; lo; hi }\n\nlet log ?(invert = false) ~lo ~hi () =\n let lo_log = Float.log10 (Float.max 1e-300 lo) in\n let hi_log = Float.log10 (Float.max 1e-300 hi) in\n let range = hi_log -. lo_log in\n let range = if range = 0. then 1. else range in\n let to_unit, from_unit =\n maybe_invert invert\n (fun v ->\n if v <= 0. then Float.nan else (Float.log10 v -. lo_log) /. range)\n (fun u -> Float.pow 10. (lo_log +. (u *. range)))\n in\n { to_unit; from_unit; lo; hi }\n\nlet sqrt ?(invert = false) ~lo ~hi () =\n let lo_s = Float.sqrt (Float.max 0. lo) in\n let hi_s = Float.sqrt (Float.max 0. hi) in\n let range = hi_s -. lo_s in\n let range = if range = 0. then 1. else range in\n let to_unit, from_unit =\n maybe_invert invert\n (fun v -> (Float.sqrt (Float.max 0. v) -. lo_s) /. range)\n (fun u ->\n let s = lo_s +. (u *. range) in\n s *. s)\n in\n { to_unit; from_unit; lo; hi }\n\nlet asinh ?(invert = false) ~lo ~hi () =\n let lo_a = Float.asinh lo in\n let hi_a = Float.asinh hi in\n let range = hi_a -. lo_a in\n let range = if range = 0. then 1. else range in\n let to_unit, from_unit =\n maybe_invert invert\n (fun v -> (Float.asinh v -. lo_a) /. range)\n (fun u ->\n let a = lo_a +. (u *. range) in\n Float.sinh a)\n in\n { to_unit; from_unit; lo; hi }\n\nlet symlog ?(invert = false) ~linthresh ~lo ~hi () =\n let transform v =\n if Float.abs v <= linthresh then v /. linthresh\n else Float.copy_sign (1. +. Float.log10 (Float.abs v /. linthresh)) v\n in\n let inv_transform v =\n if Float.abs v <= 1. then v *. linthresh\n else Float.copy_sign (linthresh *. Float.pow 10. (Float.abs v -. 1.)) v\n in\n let lo_t = transform lo in\n let hi_t = transform hi in\n let range = hi_t -. lo_t in\n let range = if range = 0. then 1. else range in\n let to_unit, from_unit =\n maybe_invert invert\n (fun v -> (transform v -. lo_t) /. range)\n (fun u -> inv_transform (lo_t +. (u *. range)))\n in\n { to_unit; from_unit; lo; hi }\n\nlet make ?(invert = false) kind ~lo ~hi () =\n match kind with\n | `Linear -> linear ~invert ~lo ~hi ()\n | `Log -> log ~invert ~lo ~hi ()\n | `Sqrt -> sqrt ~invert ~lo ~hi ()\n | `Asinh -> asinh ~invert ~lo ~hi ()\n | `Symlog linthresh -> symlog ~invert ~linthresh ~lo ~hi ()\n"
},
{
"path": "packages/hugin/lib/scale.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Data-to-unit mapping functions.\n\n {b Internal module.} Maps data-space values to the unit interval [[0, 1]]\n for linear, logarithmic, square-root, inverse-hyperbolic-sine, and\n symmetric-log scales. When [~invert] is [true], the mapping is reversed\n ([lo] maps to [1] and [hi] to [0]). *)\n\ntype t = {\n to_unit : float -> float; (** [to_unit v] maps data value [v] to [[0, 1]]. *)\n from_unit : float -> float;\n (** [from_unit u] maps unit value [u] back to data space. *)\n lo : float; (** Lower bound in data space. *)\n hi : float; (** Upper bound in data space. *)\n}\n(** The type for scales. *)\n\nval linear : ?invert:bool -> lo:float -> hi:float -> unit -> t\n(** [linear ~lo ~hi ()] is a linear scale over [[lo, hi]]. *)\n\nval log : ?invert:bool -> lo:float -> hi:float -> unit -> t\n(** [log ~lo ~hi ()] is a base-10 logarithmic scale over [[lo, hi]]. *)\n\nval sqrt : ?invert:bool -> lo:float -> hi:float -> unit -> t\n(** [sqrt ~lo ~hi ()] is a square-root scale over [[lo, hi]]. Values below zero\n are clamped. *)\n\nval asinh : ?invert:bool -> lo:float -> hi:float -> unit -> t\n(** [asinh ~lo ~hi ()] is an inverse-hyperbolic-sine scale over [[lo, hi]].\n Transitions smoothly from linear near zero to logarithmic at large absolute\n values. Handles negative values. *)\n\nval symlog :\n ?invert:bool -> linthresh:float -> lo:float -> hi:float -> unit -> t\n(** [symlog ~linthresh ~lo ~hi ()] is a symmetric logarithmic scale. Linear\n within \\[[-linthresh];[linthresh]\\], logarithmic outside. *)\n\nval make :\n ?invert:bool ->\n [ `Linear | `Log | `Sqrt | `Asinh | `Symlog of float ] ->\n lo:float ->\n hi:float ->\n unit ->\n t\n(** [make kind ~lo ~hi ()] is a scale of the given [kind]. *)\n"
},
{
"path": "packages/hugin/lib/scene.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Scene IR — resolved primitives in device pixels *)\n\ntype primitive =\n | Path of {\n points : (float * float) array;\n close : bool;\n fill : Color.t option;\n stroke : Color.t option;\n line_width : float;\n dash : float list;\n }\n | Markers of {\n points : (float * float) array;\n shape : Spec.marker;\n size : float;\n sizes : float array option;\n fill : Color.t option;\n fills : Color.t array option;\n stroke : Color.t option;\n }\n | Text of {\n x : float;\n y : float;\n content : string;\n font : Theme.font;\n color : Color.t;\n anchor : [ `Start | `Middle | `End ];\n baseline : [ `Top | `Middle | `Bottom ];\n angle : float;\n }\n | Image of { x : float; y : float; w : float; h : float; data : Nx.uint8_t }\n | Clip of {\n x : float;\n y : float;\n w : float;\n h : float;\n children : primitive list;\n }\n | Group of primitive list\n\ntype t = { width : float; height : float; primitives : primitive list }\n\nlet rec fold_primitive f acc = function\n | Group children | Clip { children; _ } ->\n List.fold_left (fold_primitive f) acc children\n | p -> f acc p\n\nlet fold f scene acc = List.fold_left (fold_primitive f) acc scene.primitives\n"
},
{
"path": "packages/hugin/lib/scene.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Scene intermediate representation.\n\n {b Internal module.} Resolved drawing primitives in device-pixel\n coordinates. All data-space concepts are gone; backends fold over these\n primitives to produce output. *)\n\n(** {1:types Types} *)\n\ntype primitive =\n | Path of {\n points : (float * float) array;\n close : bool;\n fill : Color.t option;\n stroke : Color.t option;\n line_width : float;\n dash : float list;\n }\n | Markers of {\n points : (float * float) array;\n shape : Spec.marker;\n size : float;\n sizes : float array option;\n fill : Color.t option;\n fills : Color.t array option;\n stroke : Color.t option;\n }\n | Text of {\n x : float;\n y : float;\n content : string;\n font : Theme.font;\n color : Color.t;\n anchor : [ `Start | `Middle | `End ];\n baseline : [ `Top | `Middle | `Bottom ];\n angle : float;\n }\n | Image of { x : float; y : float; w : float; h : float; data : Nx.uint8_t }\n | Clip of {\n x : float;\n y : float;\n w : float;\n h : float;\n children : primitive list;\n }\n | Group of primitive list (** The type for drawing primitives. *)\n\ntype t = { width : float; height : float; primitives : primitive list }\n(** The type for resolved scenes. *)\n\n(** {1:traversal Traversal} *)\n\nval fold : ('a -> primitive -> 'a) -> t -> 'a -> 'a\n(** [fold f scene acc] folds [f] over every leaf primitive in [scene],\n descending into {!Clip} and {!Group} nodes. *)\n"
},
{
"path": "packages/hugin/lib/spec.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype line_style = [ `Solid | `Dashed | `Dotted | `Dash_dot ]\ntype marker = Circle | Square | Triangle | Plus | Star\n\ntype legend_loc =\n | Upper_right\n | Upper_left\n | Lower_right\n | Lower_left\n | Center\n | Right\n | Upper_center\n | Lower_center\n\ntype scale = [ `Linear | `Log | `Sqrt | `Asinh | `Symlog of float ]\ntype stretch = [ `Linear | `Log | `Sqrt | `Asinh | `Power of float ]\n\ntype mark =\n | Line of {\n x : Nx.float32_t;\n y : Nx.float32_t;\n color : Color.t option;\n line_width : float option;\n line_style : line_style option;\n step : [ `Pre | `Post | `Mid ] option;\n marker : marker option;\n label : string option;\n alpha : float option;\n }\n | Point of {\n x : Nx.float32_t;\n y : Nx.float32_t;\n color : Color.t option;\n color_by : Nx.float32_t option;\n size : float option;\n size_by : Nx.float32_t option;\n marker : marker option;\n label : string option;\n alpha : float option;\n }\n | Bar of {\n x : Nx.float32_t;\n height : Nx.float32_t;\n width : float option;\n bottom : float;\n color : Color.t option;\n label : string option;\n alpha : float option;\n }\n | Hist of {\n x : Nx.float32_t;\n bins : [ `Num of int | `Edges of float array ];\n density : bool;\n color : Color.t option;\n label : string option;\n }\n | Image of {\n data : Nx.uint8_t;\n extent : (float * float * float * float) option;\n }\n | Text_mark of {\n x : float;\n y : float;\n content : string;\n color : Color.t option;\n font_size : float option;\n }\n | Hline of {\n y : float;\n color : Color.t option;\n line_width : float option;\n line_style : line_style option;\n label : string option;\n alpha : float option;\n }\n | Vline of {\n x : float;\n color : Color.t option;\n line_width : float option;\n line_style : line_style option;\n label : string option;\n alpha : float option;\n }\n | Abline of {\n slope : float;\n intercept : float;\n color : Color.t option;\n line_width : float option;\n line_style : line_style option;\n label : string option;\n alpha : float option;\n }\n | Fill_between of {\n x : Nx.float32_t;\n y1 : Nx.float32_t;\n y2 : Nx.float32_t;\n where : Nx.float32_t option;\n color : Color.t option;\n alpha : float option;\n label : string option;\n }\n | Hspan of {\n y0 : float;\n y1 : float;\n color : Color.t option;\n alpha : float option;\n label : string option;\n }\n | Vspan of {\n x0 : float;\n x1 : float;\n color : Color.t option;\n alpha : float option;\n label : string option;\n }\n | Errorbar of {\n x : Nx.float32_t;\n y : Nx.float32_t;\n yerr :\n [ `Symmetric of Nx.float32_t\n | `Asymmetric of Nx.float32_t * Nx.float32_t ];\n xerr :\n [ `Symmetric of Nx.float32_t\n | `Asymmetric of Nx.float32_t * Nx.float32_t ]\n option;\n color : Color.t option;\n line_width : float option;\n cap_size : float option;\n label : string option;\n alpha : float option;\n }\n | Heatmap of {\n data : Nx.float32_t;\n cmap : Cmap.t option;\n annotate : bool;\n vmin : float option;\n vmax : float option;\n fmt : (float -> string) option;\n }\n | Imshow of {\n data : Nx.float32_t;\n stretch : stretch;\n cmap : Cmap.t option;\n vmin : float option;\n vmax : float option;\n }\n | Contour of {\n data : Nx.float32_t;\n x0 : float;\n x1 : float;\n y0 : float;\n y1 : float;\n levels : [ `Num of int | `Values of float array ];\n filled : bool;\n cmap : Cmap.t option;\n color : Color.t option;\n line_width : float option;\n label : string option;\n alpha : float option;\n }\n\ntype decoration =\n | Title of string\n | Xlabel of string\n | Ylabel of string\n | Xlim of float * float\n | Ylim of float * float\n | Xscale of scale\n | Yscale of scale\n | Xinvert\n | Yinvert\n | Grid_visible of bool\n | Legend of legend_loc * int\n | Xticks of (float * string) list\n | Yticks of (float * string) list\n | With_theme of Theme.t\n | Xtick_format of (float -> string)\n | Ytick_format of (float -> string)\n | Frame of bool\n\ntype t =\n | Mark of mark\n | Layers of t list\n | Decorated of { inner : t; decorations : decoration list }\n | Grid of { rows : t list list; gap : float }\n\n(* Mark constructors *)\n\nlet line ~x ~y ?color ?line_width ?line_style ?step ?marker ?label ?alpha () =\n Mark\n (Line { x; y; color; line_width; line_style; step; marker; label; alpha })\n\nlet point ~x ~y ?color ?color_by ?size ?size_by ?marker ?label ?alpha () =\n Mark (Point { x; y; color; color_by; size; size_by; marker; label; alpha })\n\nlet bar ~x ~height ?width ?(bottom = 0.) ?color ?label ?alpha () =\n Mark (Bar { x; height; width; bottom; color; label; alpha })\n\nlet hist ~x ?(bins = `Num 10) ?(density = false) ?color ?label () =\n Mark (Hist { x; bins; density; color; label })\n\nlet image ?extent data = Mark (Image { data; extent })\n\nlet text ~x ~y s ?color ?font_size () =\n Mark (Text_mark { x; y; content = s; color; font_size })\n\nlet hline ~y ?color ?line_width ?line_style ?label ?alpha () =\n Mark (Hline { y; color; line_width; line_style; label; alpha })\n\nlet vline ~x ?color ?line_width ?line_style ?label ?alpha () =\n Mark (Vline { x; color; line_width; line_style; label; alpha })\n\nlet abline ~slope ~intercept ?color ?line_width ?line_style ?label ?alpha () =\n Mark\n (Abline { slope; intercept; color; line_width; line_style; label; alpha })\n\nlet fill_between ~x ~y1 ~y2 ?where ?color ?alpha ?label () =\n Mark (Fill_between { x; y1; y2; where; color; alpha; label })\n\nlet hspan ~y0 ~y1 ?color ?alpha ?label () =\n Mark (Hspan { y0; y1; color; alpha; label })\n\nlet vspan ~x0 ~x1 ?color ?alpha ?label () =\n Mark (Vspan { x0; x1; color; alpha; label })\n\nlet errorbar ~x ~y ~yerr ?xerr ?color ?line_width ?cap_size ?label ?alpha () =\n Mark\n (Errorbar { x; y; yerr; xerr; color; line_width; cap_size; label; alpha })\n\nlet heatmap ~data ?(annotate = false) ?cmap ?vmin ?vmax ?fmt () =\n Mark (Heatmap { data; cmap; annotate; vmin; vmax; fmt })\n\nlet imshow ~data ?(stretch = `Linear) ?cmap ?vmin ?vmax () =\n Mark (Imshow { data; stretch; cmap; vmin; vmax })\n\nlet contour ~data ~x0 ~x1 ~y0 ~y1 ?(levels = `Num 8) ?(filled = false) ?cmap\n ?color ?line_width ?label ?alpha () =\n Mark\n (Contour\n {\n data;\n x0;\n x1;\n y0;\n y1;\n levels;\n filled;\n cmap;\n color;\n line_width;\n label;\n alpha;\n })\n\n(* Composition *)\n\nlet layers ts = Layers ts\n\n(* Decorations *)\n\nlet decorate d = function\n | Decorated r -> Decorated { r with decorations = d :: r.decorations }\n | t -> Decorated { inner = t; decorations = [ d ] }\n\nlet title s t = decorate (Title s) t\nlet xlabel s t = decorate (Xlabel s) t\nlet ylabel s t = decorate (Ylabel s) t\nlet xlim lo hi t = decorate (Xlim (lo, hi)) t\nlet ylim lo hi t = decorate (Ylim (lo, hi)) t\nlet xscale s t = decorate (Xscale s) t\nlet yscale s t = decorate (Yscale s) t\nlet xinvert t = decorate Xinvert t\nlet yinvert t = decorate Yinvert t\nlet grid_lines visible t = decorate (Grid_visible visible) t\nlet legend ?(loc = Upper_right) ?(ncol = 1) t = decorate (Legend (loc, ncol)) t\nlet xticks ticks t = decorate (Xticks ticks) t\nlet yticks ticks t = decorate (Yticks ticks) t\nlet with_theme th t = decorate (With_theme th) t\nlet xtick_format fmt t = decorate (Xtick_format fmt) t\nlet ytick_format fmt t = decorate (Ytick_format fmt) t\nlet frame v t = decorate (Frame v) t\n\nlet no_axes t =\n t |> decorate (Frame false) |> decorate (Xticks []) |> decorate (Yticks [])\n |> decorate (Grid_visible false)\n\n(* Layout *)\n\nlet grid_layout ?(gap = 0.05) rows = Grid { rows; gap }\n"
},
{
"path": "packages/hugin/lib/spec.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Immutable plot specifications.\n\n {b Internal module.} The specification tree is the user-facing\n representation of a plot. {!Prepared.compile} resolves data-dependent work;\n {!Resolve} turns the result into a {!Scene.t}. *)\n\n(** {1:types Types} *)\n\ntype line_style = [ `Solid | `Dashed | `Dotted | `Dash_dot ]\n(** The type for line dash patterns. *)\n\ntype marker =\n | Circle\n | Square\n | Triangle\n | Plus\n | Star (** The type for point marker shapes. *)\n\ntype legend_loc =\n | Upper_right\n | Upper_left\n | Lower_right\n | Lower_left\n | Center\n | Right\n | Upper_center\n | Lower_center (** The type for legend placement. *)\n\ntype scale = [ `Linear | `Log | `Sqrt | `Asinh | `Symlog of float ]\n(** The type for axis scales. [`Sqrt] and [`Asinh] handle zero gracefully.\n [`Symlog linthresh] is linear within \\[[-linthresh];[linthresh]\\] and\n logarithmic outside. *)\n\ntype stretch = [ `Linear | `Log | `Sqrt | `Asinh | `Power of float ]\n(** The type for image stretch functions. [`Power a] raises normalized values to\n the power [a]. *)\n\ntype mark =\n | Line of {\n x : Nx.float32_t;\n y : Nx.float32_t;\n color : Color.t option;\n line_width : float option;\n line_style : line_style option;\n step : [ `Pre | `Post | `Mid ] option;\n marker : marker option;\n label : string option;\n alpha : float option;\n }\n | Point of {\n x : Nx.float32_t;\n y : Nx.float32_t;\n color : Color.t option;\n color_by : Nx.float32_t option;\n size : float option;\n size_by : Nx.float32_t option;\n marker : marker option;\n label : string option;\n alpha : float option;\n }\n | Bar of {\n x : Nx.float32_t;\n height : Nx.float32_t;\n width : float option;\n bottom : float;\n color : Color.t option;\n label : string option;\n alpha : float option;\n }\n | Hist of {\n x : Nx.float32_t;\n bins : [ `Num of int | `Edges of float array ];\n density : bool;\n color : Color.t option;\n label : string option;\n }\n | Image of {\n data : Nx.uint8_t;\n extent : (float * float * float * float) option;\n }\n | Text_mark of {\n x : float;\n y : float;\n content : string;\n color : Color.t option;\n font_size : float option;\n }\n | Hline of {\n y : float;\n color : Color.t option;\n line_width : float option;\n line_style : line_style option;\n label : string option;\n alpha : float option;\n }\n | Vline of {\n x : float;\n color : Color.t option;\n line_width : float option;\n line_style : line_style option;\n label : string option;\n alpha : float option;\n }\n | Abline of {\n slope : float;\n intercept : float;\n color : Color.t option;\n line_width : float option;\n line_style : line_style option;\n label : string option;\n alpha : float option;\n }\n | Fill_between of {\n x : Nx.float32_t;\n y1 : Nx.float32_t;\n y2 : Nx.float32_t;\n where : Nx.float32_t option;\n color : Color.t option;\n alpha : float option;\n label : string option;\n }\n | Hspan of {\n y0 : float;\n y1 : float;\n color : Color.t option;\n alpha : float option;\n label : string option;\n }\n | Vspan of {\n x0 : float;\n x1 : float;\n color : Color.t option;\n alpha : float option;\n label : string option;\n }\n | Errorbar of {\n x : Nx.float32_t;\n y : Nx.float32_t;\n yerr :\n [ `Symmetric of Nx.float32_t\n | `Asymmetric of Nx.float32_t * Nx.float32_t ];\n xerr :\n [ `Symmetric of Nx.float32_t\n | `Asymmetric of Nx.float32_t * Nx.float32_t ]\n option;\n color : Color.t option;\n line_width : float option;\n cap_size : float option;\n label : string option;\n alpha : float option;\n }\n | Heatmap of {\n data : Nx.float32_t;\n cmap : Cmap.t option;\n annotate : bool;\n vmin : float option;\n vmax : float option;\n fmt : (float -> string) option;\n }\n | Imshow of {\n data : Nx.float32_t;\n stretch : stretch;\n cmap : Cmap.t option;\n vmin : float option;\n vmax : float option;\n }\n | Contour of {\n data : Nx.float32_t;\n x0 : float;\n x1 : float;\n y0 : float;\n y1 : float;\n levels : [ `Num of int | `Values of float array ];\n filled : bool;\n cmap : Cmap.t option;\n color : Color.t option;\n line_width : float option;\n label : string option;\n alpha : float option;\n }\n (** The type for visual marks. Each constructor carries the data arrays\n and visual properties for one layer. *)\n\ntype decoration =\n | Title of string\n | Xlabel of string\n | Ylabel of string\n | Xlim of float * float\n | Ylim of float * float\n | Xscale of scale\n | Yscale of scale\n | Xinvert\n | Yinvert\n | Grid_visible of bool\n | Legend of legend_loc * int\n | Xticks of (float * string) list\n | Yticks of (float * string) list\n | With_theme of Theme.t\n | Xtick_format of (float -> string)\n | Ytick_format of (float -> string)\n | Frame of bool\n (** The type for plot decorations. Applied via {!Decorated} nodes. *)\n\ntype t =\n | Mark of mark\n | Layers of t list\n | Decorated of { inner : t; decorations : decoration list }\n | Grid of { rows : t list list; gap : float }\n (** The type for plot specifications. An immutable tree composed via mark\n constructors, {!layers}, decoration functions, and {!grid_layout}. *)\n\n(** {1:marks Mark constructors} *)\n\nval line :\n x:Nx.float32_t ->\n y:Nx.float32_t ->\n ?color:Color.t ->\n ?line_width:float ->\n ?line_style:line_style ->\n ?step:[ `Pre | `Post | `Mid ] ->\n ?marker:marker ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [line ~x ~y ()] is a line mark. *)\n\nval point :\n x:Nx.float32_t ->\n y:Nx.float32_t ->\n ?color:Color.t ->\n ?color_by:Nx.float32_t ->\n ?size:float ->\n ?size_by:Nx.float32_t ->\n ?marker:marker ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [point ~x ~y ()] is a scatter mark. *)\n\nval bar :\n x:Nx.float32_t ->\n height:Nx.float32_t ->\n ?width:float ->\n ?bottom:float ->\n ?color:Color.t ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [bar ~x ~height ()] is a bar mark. [bottom] defaults to [0.]. *)\n\nval hist :\n x:Nx.float32_t ->\n ?bins:[ `Num of int | `Edges of float array ] ->\n ?density:bool ->\n ?color:Color.t ->\n ?label:string ->\n unit ->\n t\n(** [hist ~x ()] is a histogram mark. [bins] defaults to [`Num 10]. *)\n\nval image : ?extent:float * float * float * float -> Nx.uint8_t -> t\n(** [image ?extent data] is an image mark. When [extent] is\n [(xmin, xmax, ymin, ymax)], the image is placed in data coordinates. *)\n\nval text :\n x:float ->\n y:float ->\n string ->\n ?color:Color.t ->\n ?font_size:float ->\n unit ->\n t\n(** [text ~x ~y s ()] is a text mark at [(x, y)]. *)\n\nval hline :\n y:float ->\n ?color:Color.t ->\n ?line_width:float ->\n ?line_style:line_style ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [hline ~y ()] is a horizontal reference line. *)\n\nval vline :\n x:float ->\n ?color:Color.t ->\n ?line_width:float ->\n ?line_style:line_style ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [vline ~x ()] is a vertical reference line. *)\n\nval abline :\n slope:float ->\n intercept:float ->\n ?color:Color.t ->\n ?line_width:float ->\n ?line_style:line_style ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [abline ~slope ~intercept ()] is a diagonal line [y = slope * x + intercept]\n spanning the full plot area. *)\n\nval fill_between :\n x:Nx.float32_t ->\n y1:Nx.float32_t ->\n y2:Nx.float32_t ->\n ?where:Nx.float32_t ->\n ?color:Color.t ->\n ?alpha:float ->\n ?label:string ->\n unit ->\n t\n(** [fill_between ~x ~y1 ~y2 ()] is a filled area between two curves. [where] is\n a mask array: only fill where [where.(i) > 0.]. *)\n\nval hspan :\n y0:float ->\n y1:float ->\n ?color:Color.t ->\n ?alpha:float ->\n ?label:string ->\n unit ->\n t\n(** [hspan ~y0 ~y1 ()] is a horizontal shaded band. *)\n\nval vspan :\n x0:float ->\n x1:float ->\n ?color:Color.t ->\n ?alpha:float ->\n ?label:string ->\n unit ->\n t\n(** [vspan ~x0 ~x1 ()] is a vertical shaded band. *)\n\nval errorbar :\n x:Nx.float32_t ->\n y:Nx.float32_t ->\n yerr:\n [ `Symmetric of Nx.float32_t | `Asymmetric of Nx.float32_t * Nx.float32_t ] ->\n ?xerr:\n [ `Symmetric of Nx.float32_t | `Asymmetric of Nx.float32_t * Nx.float32_t ] ->\n ?color:Color.t ->\n ?line_width:float ->\n ?cap_size:float ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [errorbar ~x ~y ~yerr ()] is an error bar mark. *)\n\nval heatmap :\n data:Nx.float32_t ->\n ?annotate:bool ->\n ?cmap:Cmap.t ->\n ?vmin:float ->\n ?vmax:float ->\n ?fmt:(float -> string) ->\n unit ->\n t\n(** [heatmap ~data ()] is a heatmap mark. [data] has shape [[|rows; cols|]]. *)\n\nval imshow :\n data:Nx.float32_t ->\n ?stretch:stretch ->\n ?cmap:Cmap.t ->\n ?vmin:float ->\n ?vmax:float ->\n unit ->\n t\n(** [imshow ~data ()] is a colormapped image mark. [stretch] defaults to\n [`Linear]. *)\n\nval contour :\n data:Nx.float32_t ->\n x0:float ->\n x1:float ->\n y0:float ->\n y1:float ->\n ?levels:[ `Num of int | `Values of float array ] ->\n ?filled:bool ->\n ?cmap:Cmap.t ->\n ?color:Color.t ->\n ?line_width:float ->\n ?label:string ->\n ?alpha:float ->\n unit ->\n t\n(** [contour ~data ~x0 ~x1 ~y0 ~y1 ()] is a contour mark. [levels] defaults to\n [`Num 8]. [filled] defaults to [false]. *)\n\n(** {1:composition Composition} *)\n\nval layers : t list -> t\n(** [layers marks] overlays [marks] on shared axes. *)\n\n(** {1:decorations Decorations} *)\n\nval title : string -> t -> t\n(** [title s t] adds plot title [s]. *)\n\nval xlabel : string -> t -> t\n(** [xlabel s t] adds x-axis label [s]. *)\n\nval ylabel : string -> t -> t\n(** [ylabel s t] adds y-axis label [s]. *)\n\nval xlim : float -> float -> t -> t\n(** [xlim lo hi t] fixes the x-axis range. *)\n\nval ylim : float -> float -> t -> t\n(** [ylim lo hi t] fixes the y-axis range. *)\n\nval xscale : scale -> t -> t\n(** [xscale s t] sets the x-axis scale. *)\n\nval yscale : scale -> t -> t\n(** [yscale s t] sets the y-axis scale. *)\n\nval xinvert : t -> t\n(** [xinvert t] inverts the x-axis direction (values increase right-to-left). *)\n\nval yinvert : t -> t\n(** [yinvert t] inverts the y-axis direction (values increase top-to-bottom). *)\n\nval grid_lines : bool -> t -> t\n(** [grid_lines visible t] shows or hides grid lines. *)\n\nval legend : ?loc:legend_loc -> ?ncol:int -> t -> t\n(** [legend t] shows the legend. [loc] defaults to {!Upper_right}. [ncol]\n defaults to [1]; set higher for multi-column layouts. *)\n\nval xticks : (float * string) list -> t -> t\n(** [xticks ticks t] sets explicit x-axis tick positions and labels. *)\n\nval yticks : (float * string) list -> t -> t\n(** [yticks ticks t] sets explicit y-axis tick positions and labels. *)\n\nval with_theme : Theme.t -> t -> t\n(** [with_theme th t] overrides the rendering theme. *)\n\nval xtick_format : (float -> string) -> t -> t\n(** [xtick_format fmt t] formats x-axis tick labels with [fmt]. *)\n\nval ytick_format : (float -> string) -> t -> t\n(** [ytick_format fmt t] formats y-axis tick labels with [fmt]. *)\n\nval frame : bool -> t -> t\n(** [frame visible t] shows or hides the axis border rectangle. *)\n\nval no_axes : t -> t\n(** [no_axes t] hides the axis frame, ticks, and tick labels. The full panel\n area is used for marks. Title is preserved. Useful for image grids. *)\n\n(** {1:layout Layout} *)\n\nval grid_layout : ?gap:float -> t list list -> t\n(** [grid_layout rows] arranges specs in a grid. [gap] defaults to [0.05]. *)\n"
},
{
"path": "packages/hugin/lib/svg_backend.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* SVG backend *)\n\n(* Text measurer *)\n\nlet text_measurer ~(font : Theme.font) s =\n let w = float (String.length s) *. font.size *. 0.6 in\n let h = font.size in\n (w, h)\n\n(* Helpers *)\n\nlet color_to_rgb_string c =\n let r, g, b, _ = Color.to_rgba c in\n Printf.sprintf \"rgb(%d,%d,%d)\"\n (Float.to_int (r *. 255.))\n (Float.to_int (g *. 255.))\n (Float.to_int (b *. 255.))\n\nlet add_fill buf = function\n | None -> Buffer.add_string buf \" fill=\\\"none\\\"\"\n | Some c ->\n Printf.bprintf buf \" fill=\\\"%s\\\"\" (color_to_rgb_string c);\n let a = Color.alpha c in\n if a < 1. then Printf.bprintf buf \" fill-opacity=\\\"%.3g\\\"\" a\n\nlet add_stroke buf = function\n | None -> Buffer.add_string buf \" stroke=\\\"none\\\"\"\n | Some c ->\n Printf.bprintf buf \" stroke=\\\"%s\\\"\" (color_to_rgb_string c);\n let a = Color.alpha c in\n if a < 1. then Printf.bprintf buf \" stroke-opacity=\\\"%.3g\\\"\" a\n\nlet text_anchor_string = function\n | `Start -> \"start\"\n | `Middle -> \"middle\"\n | `End -> \"end\"\n\nlet dominant_baseline_string = function\n | `Top -> \"text-before-edge\"\n | `Middle -> \"central\"\n | `Bottom -> \"text-after-edge\"\n\nlet escape_xml s =\n let buf = Buffer.create (String.length s) in\n String.iter\n (function\n | '<' -> Buffer.add_string buf \"<\"\n | '>' -> Buffer.add_string buf \">\"\n | '&' -> Buffer.add_string buf \"&\"\n | '\"' -> Buffer.add_string buf \""\"\n | c -> Buffer.add_char buf c)\n s;\n Buffer.contents buf\n\n(* Marker shapes *)\n\nlet marker_path shape size =\n let hs = size /. 2. in\n match shape with\n | Spec.Circle ->\n Printf.sprintf \"M %g 0 A %g %g 0 1 1 %g 0 A %g %g 0 1 1 %g 0 Z\" (-.hs) hs\n hs hs hs hs (-.hs)\n | Spec.Square ->\n Printf.sprintf \"M %g %g L %g %g L %g %g L %g %g Z\" (-.hs) (-.hs) hs (-.hs)\n hs hs (-.hs) hs\n | Spec.Triangle ->\n Printf.sprintf \"M 0 %g L %g %g L %g %g Z\" (-.hs) hs hs (-.hs) hs\n | Spec.Plus ->\n Printf.sprintf \"M %g 0 L %g 0 M 0 %g L 0 %g\" (-.hs) hs (-.hs) hs\n | Spec.Star ->\n let d = hs *. 0.707 in\n Printf.sprintf\n \"M %g 0 L %g 0 M 0 %g L 0 %g M %g %g L %g %g M %g %g L %g %g\" (-.hs) hs\n (-.hs) hs (-.d) (-.d) d d d (-.d) (-.d) d\n\n(* Primitive rendering — ids threaded through to avoid global state *)\n\ntype ids = { mutable clip_id : int; mutable marker_id : int }\n\nlet fresh_clip ids =\n ids.clip_id <- ids.clip_id + 1;\n Printf.sprintf \"clip-%d\" ids.clip_id\n\nlet fresh_marker ids =\n ids.marker_id <- ids.marker_id + 1;\n Printf.sprintf \"marker-%d\" ids.marker_id\n\nlet rec render_primitive ids buf = function\n | Scene.Path { points; close; fill; stroke; line_width; dash } ->\n if Array.length points < 2 then ()\n else begin\n Buffer.add_string buf \"\n if i = 0 then Printf.bprintf buf \"M %g %g\" x y\n else Printf.bprintf buf \" L %g %g\" x y)\n points;\n if close then Buffer.add_string buf \" Z\";\n Buffer.add_char buf '\"';\n add_fill buf fill;\n add_stroke buf stroke;\n if line_width > 0. then\n Printf.bprintf buf \" stroke-width=\\\"%g\\\"\" line_width;\n begin match dash with\n | [] -> ()\n | ds ->\n Buffer.add_string buf \" stroke-dasharray=\\\"\";\n List.iteri\n (fun i d ->\n if i > 0 then Buffer.add_char buf ',';\n Printf.bprintf buf \"%g\" d)\n ds;\n Buffer.add_char buf '\"'\n end;\n Buffer.add_string buf \"/>\\n\"\n end\n | Scene.Markers { points; shape; size; sizes; fill; fills; stroke } ->\n let stroke_only =\n match shape with Spec.Plus | Spec.Star -> true | _ -> false\n in\n begin match (fills, sizes) with\n | None, None ->\n let id = fresh_marker ids in\n let d = marker_path shape size in\n Printf.bprintf buf \" fill | None -> stroke in\n add_stroke buf stroke_c;\n Printf.bprintf buf \" stroke-width=\\\"%g\\\"\"\n (Float.max 1. (size *. 0.15))\n end\n else begin\n add_fill buf fill;\n add_stroke buf stroke;\n if stroke <> None then Buffer.add_string buf \" stroke-width=\\\"1\\\"\"\n end;\n Buffer.add_string buf \"/> \\n\" id x\n y)\n points\n | _ ->\n Array.iteri\n (fun i (x, y) ->\n let s = match sizes with Some ss -> ss.(i) | None -> size in\n let f =\n match fills with Some fs -> Some fs.(i) | None -> fill\n in\n let d = marker_path shape s in\n Printf.bprintf buf \" f | None -> stroke in\n add_stroke buf stroke_c;\n Printf.bprintf buf \" stroke-width=\\\"%g\\\"\"\n (Float.max 1. (s *. 0.15))\n end\n else begin\n add_fill buf f;\n add_stroke buf stroke;\n if stroke <> None then\n Buffer.add_string buf \" stroke-width=\\\"1\\\"\"\n end;\n Buffer.add_string buf \"/>\\n\")\n points\n end\n | Scene.Text { x; y; content; font; color; anchor; baseline; angle } ->\n Printf.bprintf buf \" Buffer.add_string buf \" font-weight=\\\"bold\\\"\"\n | `Normal -> ()\n end;\n Printf.bprintf buf \" fill=\\\"%s\\\"\" (color_to_rgb_string color);\n let a = Color.alpha color in\n if a < 1. then Printf.bprintf buf \" fill-opacity=\\\"%.3g\\\"\" a;\n Printf.bprintf buf \" text-anchor=\\\"%s\\\"\" (text_anchor_string anchor);\n Printf.bprintf buf \" dominant-baseline=\\\"%s\\\"\"\n (dominant_baseline_string baseline);\n if angle <> 0. then\n Printf.bprintf buf \" transform=\\\"rotate(%g,%g,%g)\\\"\"\n (angle *. -180. /. Float.pi)\n x y;\n Printf.bprintf buf \">%s \\n\" (escape_xml content)\n | Scene.Image { x; y; w; h; data } ->\n let b64 = Image_util.nx_to_png_base64 data in\n Printf.bprintf buf \"\\n\" b64\n | Scene.Clip { x; y; w; h; children } ->\n let id = fresh_clip ids in\n Printf.bprintf buf\n \"\\n\" id;\n List.iter (render_primitive ids buf) children;\n Buffer.add_string buf \" \\n\"\n | Scene.Group children ->\n Buffer.add_string buf \"\\n\";\n List.iter (render_primitive ids buf) children;\n Buffer.add_string buf \" \\n\"\n\n(* Entry points *)\n\nlet render (scene : Scene.t) =\n let ids = { clip_id = 0; marker_id = 0 } in\n let buf = Buffer.create 4096 in\n Printf.bprintf buf \"\\n\";\n Printf.bprintf buf\n \"\\n\"\n scene.width scene.height scene.width scene.height;\n List.iter (render_primitive ids buf) scene.primitives;\n Buffer.add_string buf \" \\n\";\n Buffer.contents buf\n\nlet render_to_file filename scene =\n let s = render scene in\n let oc = open_out filename in\n output_string oc s;\n close_out oc\n"
},
{
"path": "packages/hugin/lib/svg_backend.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** SVG rendering backend.\n\n {b Internal module.} Renders {!Scene.t} to SVG markup. Pure OCaml, no\n external dependencies beyond Cairo for image encoding. *)\n\n(** {1:measurer Text measurement} *)\n\nval text_measurer : Resolve.text_measurer\n(** [text_measurer] estimates text dimensions from character count and font\n size. Heuristic: width is [String.length s * 0.6 * font.size]. *)\n\n(** {1:rendering Rendering} *)\n\nval render : Scene.t -> string\n(** [render scene] is [scene] as an SVG document string. *)\n\nval render_to_file : string -> Scene.t -> unit\n(** [render_to_file filename scene] writes [scene] as an SVG file. *)\n"
},
{
"path": "packages/hugin/lib/theme.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype font = { family : string; size : float; weight : [ `Normal | `Bold ] }\ntype line = { color : Color.t; width : float; dash : float list }\n\ntype t = {\n background : Color.t;\n palette : Color.t array;\n sequential : Cmap.t;\n diverging : Cmap.t;\n font_title : font;\n font_label : font;\n font_tick : font;\n axis : line;\n grid : line option;\n tick_length : float;\n padding : float;\n title_gap : float;\n label_gap : float;\n scale_factor : float;\n line_width : float;\n marker_size : float;\n}\n\nlet axis_color = Color.oklch ~l:0.3 ~c:0. ~h:0. ()\nlet grid_color = Color.with_alpha 0.15 axis_color\n\nlet default =\n {\n background = Color.oklch ~l:0.985 ~c:0. ~h:0. ();\n palette =\n [|\n Color.orange;\n Color.sky_blue;\n Color.green;\n Color.darken 0.1 Color.yellow;\n Color.blue;\n Color.vermillion;\n Color.purple;\n Color.black;\n |];\n sequential = Cmap.viridis;\n diverging = Cmap.coolwarm;\n font_title = { family = \"sans-serif\"; size = 28.; weight = `Bold };\n font_label = { family = \"sans-serif\"; size = 22.; weight = `Normal };\n font_tick = { family = \"sans-serif\"; size = 18.; weight = `Normal };\n axis = { color = axis_color; width = 2.; dash = [] };\n grid = Some { color = grid_color; width = 1.; dash = [] };\n tick_length = 10.;\n padding = 24.;\n title_gap = 16.;\n label_gap = 12.;\n scale_factor = 1.;\n line_width = 3.;\n marker_size = 10.;\n }\n\nlet dark_bg = Color.oklch ~l:0.15 ~c:0. ~h:0. ()\nlet dark_fg = Color.oklch ~l:0.8 ~c:0. ~h:0. ()\nlet dark_grid = Color.with_alpha 0.2 dark_fg\n\nlet dark =\n {\n default with\n background = dark_bg;\n palette =\n [|\n Color.orange;\n Color.sky_blue;\n Color.green;\n Color.yellow;\n Color.blue;\n Color.vermillion;\n Color.purple;\n Color.white;\n |];\n axis = { color = dark_fg; width = 2.; dash = [] };\n grid = Some { color = dark_grid; width = 1.; dash = [] };\n }\n\nlet minimal =\n { default with axis = { default.axis with width = 1. }; grid = None }\n\nlet paper t = { t with scale_factor = 1.0 }\nlet notebook t = { t with scale_factor = 1.3 }\nlet talk t = { t with scale_factor = 1.6 }\nlet poster t = { t with scale_factor = 2.0 }\n"
},
{
"path": "packages/hugin/lib/theme.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Visual themes.\n\n A theme controls every non-data visual element: background color,\n typography, axes, grid, spacing, and default data palettes.\n\n Themes separate two orthogonal concerns: {e style} (aesthetic appearance)\n and {e context} (scaling for the output medium). The {!paper}, {!notebook},\n {!talk}, and {!poster} functions adjust {!field-scale_factor} to uniformly\n scale all visual elements for the target medium. *)\n\n(** {1:types Types} *)\n\ntype font = { family : string; size : float; weight : [ `Normal | `Bold ] }\n(** The type for font specifications. [size] is in points before\n {!field-scale_factor} is applied. *)\n\ntype line = { color : Color.t; width : float; dash : float list }\n(** The type for line styles. [dash] is a list of on/off lengths; empty means\n solid. *)\n\ntype t = {\n background : Color.t;\n palette : Color.t array;\n sequential : Cmap.t;\n diverging : Cmap.t;\n font_title : font;\n font_label : font;\n font_tick : font;\n axis : line;\n grid : line option;\n tick_length : float;\n padding : float;\n title_gap : float;\n label_gap : float;\n scale_factor : float;\n line_width : float;\n marker_size : float;\n}\n(** The type for themes. All dimensional values (font sizes, line widths, gaps)\n are multiplied by {!field-scale_factor} at render time. *)\n\n(** {1:predefined Predefined themes} *)\n\nval default : t\n(** [default] is a light theme with subtle grid, Okabe-Ito categorical palette,\n and Tufte-informed defaults. *)\n\nval dark : t\n(** [dark] is a dark-background theme. *)\n\nval minimal : t\n(** [minimal] is a theme with no grid and thin axes. *)\n\n(** {1:context Context scaling} *)\n\nval paper : t -> t\n(** [paper t] is [t] with [scale_factor = 1.0]. *)\n\nval notebook : t -> t\n(** [notebook t] is [t] with [scale_factor = 1.3]. *)\n\nval talk : t -> t\n(** [talk t] is [t] with [scale_factor = 1.6]. *)\n\nval poster : t -> t\n(** [poster t] is [t] with [scale_factor = 2.0]. *)\n"
},
{
"path": "packages/hugin/lib/ticks.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Nice tick generation *)\n\nlet nice_num v round =\n let exp = Float.floor (Float.log10 v) in\n let frac = v /. Float.pow 10. exp in\n let nice =\n if round then\n begin if frac < 1.5 then 1.\n else if frac < 3. then 2.\n else if frac < 7. then 5.\n else 10.\n end\n else\n begin if frac <= 1. then 1.\n else if frac <= 2. then 2.\n else if frac <= 5. then 5.\n else 10.\n end\n in\n nice *. Float.pow 10. exp\n\nlet format_tick v =\n let v = if Float.abs v < 1e-14 then 0. else v in\n Printf.sprintf \"%.6g\" v\n\nlet generate_linear ~lo ~hi ~max_ticks =\n let range = nice_num (hi -. lo) false in\n let step = nice_num (range /. float max_ticks) true in\n let lo' = Float.floor (lo /. step) *. step in\n let acc = ref [] in\n let v = ref lo' in\n while !v <= hi +. (step *. 0.5) do\n if !v >= lo -. (step *. 0.001) && !v <= hi +. (step *. 0.001) then\n acc := (!v, format_tick !v) :: !acc;\n v := !v +. step\n done;\n List.rev !acc\n\nlet format_log_tick e =\n let ei = int_of_float e in\n if ei = 0 then \"1\" else if ei = 1 then \"10\" else Printf.sprintf \"10^%d\" ei\n\nlet generate_log ~lo ~hi ~max_ticks =\n let lo_exp = Float.floor (Float.log10 (Float.max 1e-300 lo)) in\n let hi_exp = Float.ceil (Float.log10 (Float.max 1e-300 hi)) in\n let n_decades = int_of_float (hi_exp -. lo_exp) in\n let stride = Float.of_int (max 1 ((n_decades + max_ticks - 1) / max_ticks)) in\n let acc = ref [] in\n let e = ref lo_exp in\n while !e <= hi_exp do\n let v = Float.pow 10. !e in\n if v >= lo *. 0.999 && v <= hi *. 1.001 then\n acc := (v, format_log_tick !e) :: !acc;\n e := !e +. stride\n done;\n List.rev !acc\n\n(* Sqrt ticks: generate in data space using nice linear ticks *)\n\nlet generate_sqrt ~lo ~hi ~max_ticks =\n let lo = Float.max 0. lo in\n generate_linear ~lo ~hi ~max_ticks\n\n(* Asinh ticks: pick nice values in data space *)\n\nlet generate_asinh ~lo ~hi ~max_ticks =\n if lo >= 0. then generate_linear ~lo ~hi ~max_ticks\n else generate_linear ~lo ~hi ~max_ticks\n\n(* Symlog ticks: linear ticks inside linthresh, log ticks outside *)\n\nlet generate_symlog ~linthresh ~lo ~hi ~max_ticks =\n let ticks = ref [] in\n (* Linear region *)\n let lin_lo = Float.max lo (-.linthresh) in\n let lin_hi = Float.min hi linthresh in\n if lin_lo < lin_hi then begin\n let lin_ticks =\n generate_linear ~lo:lin_lo ~hi:lin_hi ~max_ticks:(max_ticks / 2)\n in\n ticks := lin_ticks\n end;\n (* Positive log region *)\n if hi > linthresh then begin\n let pos_lo = Float.max linthresh lo in\n let pos_ticks = generate_log ~lo:pos_lo ~hi ~max_ticks:(max_ticks / 3) in\n ticks := !ticks @ pos_ticks\n end;\n (* Negative log region *)\n if lo < -.linthresh then begin\n let neg_hi = Float.min (-.linthresh) hi in\n let neg_lo_abs = Float.abs lo in\n let neg_hi_abs = Float.abs neg_hi in\n let pos_ticks =\n generate_log ~lo:neg_hi_abs ~hi:neg_lo_abs ~max_ticks:(max_ticks / 3)\n in\n let neg_ticks =\n List.rev_map (fun (v, _) -> (-.v, format_tick (-.v))) pos_ticks\n in\n ticks := neg_ticks @ !ticks\n end;\n !ticks\n\nlet generate kind ~lo ~hi ?(max_ticks = 8) () =\n match kind with\n | `Linear -> generate_linear ~lo ~hi ~max_ticks\n | `Log -> generate_log ~lo ~hi ~max_ticks\n | `Sqrt -> generate_sqrt ~lo ~hi ~max_ticks\n | `Asinh -> generate_asinh ~lo ~hi ~max_ticks\n | `Symlog linthresh -> generate_symlog ~linthresh ~lo ~hi ~max_ticks\n"
},
{
"path": "packages/hugin/lib/ticks.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Tick generation for axes.\n\n {b Internal module.} Produces nicely-spaced tick positions and formatted\n labels for linear, logarithmic, square-root, asinh, and symlog scales. *)\n\nval generate :\n [ `Linear | `Log | `Sqrt | `Asinh | `Symlog of float ] ->\n lo:float ->\n hi:float ->\n ?max_ticks:int ->\n unit ->\n (float * string) list\n(** [generate kind ~lo ~hi ()] is a list of [(value, label)] pairs for ticks\n spanning [[lo, hi]]. [max_ticks] defaults to [8]. *)\n"
},
{
"path": "packages/hugin/test/dune",
"content": "(tests\n (names\n test_color\n test_cmap\n test_scale\n test_ticks\n test_resolve\n test_svg_backend\n test_image_util)\n (package hugin)\n (libraries hugin nx windtrap))\n"
},
{
"path": "packages/hugin/test/test_cmap.ml",
"content": "open Hugin\nopen Windtrap\n\nlet check_float msg = equal ~msg (float 0.01)\nlet bw = Cmap.of_colors [| Color.black; Color.white |]\n\n(* eval *)\n\nlet test_eval_at_zero () =\n let c = Cmap.eval bw 0.0 in\n check_float \"lightness at 0\" (Color.lightness Color.black) (Color.lightness c)\n\nlet test_eval_at_one () =\n let c = Cmap.eval bw 1.0 in\n check_float \"lightness at 1\" (Color.lightness Color.white) (Color.lightness c)\n\nlet test_eval_negative_clamped () =\n let c0 = Cmap.eval bw 0.0 in\n let cn = Cmap.eval bw (-0.5) in\n check_float \"negative clamped\" (Color.lightness c0) (Color.lightness cn)\n\nlet test_eval_above_one_clamped () =\n let c1 = Cmap.eval bw 1.0 in\n let ch = Cmap.eval bw 1.5 in\n check_float \"above 1 clamped\" (Color.lightness c1) (Color.lightness ch)\n\n(* of_colors *)\n\nlet test_two_stops_midpoint () =\n let mid = Cmap.eval bw 0.5 in\n let l = Color.lightness mid in\n is_true ~msg:\"midpoint lightness near 0.5\" (l > 0.4 && l < 0.6)\n\nlet test_three_stops_midpoint () =\n let red = Color.rgb ~r:1. ~g:0. ~b:0. () in\n let green = Color.rgb ~r:0. ~g:1. ~b:0. () in\n let blue = Color.rgb ~r:0. ~g:0. ~b:1. () in\n let cm = Cmap.of_colors [| red; green; blue |] in\n let mid = Cmap.eval cm 0.5 in\n let r, g, _, _ = Color.to_rgba mid in\n (* At 0.5, should be near the green stop *)\n is_true ~msg:\"green channel high at midpoint\" (g > r)\n\nlet test_one_stop_raises () =\n raises_invalid_arg \"Cmap.of_colors: need at least 2 stops\" (fun () ->\n Cmap.of_colors [| Color.black |])\n\nlet test_empty_raises () =\n raises_invalid_arg \"Cmap.of_colors: need at least 2 stops\" (fun () ->\n Cmap.of_colors [||])\n\n(* predefined *)\n\nlet test_predefined_no_raise () =\n let cmaps =\n [\n Cmap.viridis;\n Cmap.plasma;\n Cmap.inferno;\n Cmap.magma;\n Cmap.cividis;\n Cmap.coolwarm;\n ]\n in\n List.iter\n (fun cm ->\n let _ = Cmap.eval cm 0.0 in\n let _ = Cmap.eval cm 0.5 in\n let _ = Cmap.eval cm 1.0 in\n ())\n cmaps\n\nlet test_viridis_endpoints () =\n let r0, _, b0, _ = Cmap.eval Cmap.viridis 0.0 |> Color.to_rgba in\n let r1, _, _, _ = Cmap.eval Cmap.viridis 1.0 |> Color.to_rgba in\n (* Viridis starts dark purple (low r, high b), ends yellow (high r) *)\n is_true ~msg:\"viridis start is dark\" (r0 < 0.4);\n is_true ~msg:\"viridis start has blue\" (b0 > 0.2);\n is_true ~msg:\"viridis end is bright\" (r1 > 0.8)\n\nlet () =\n run \"Cmap\"\n [\n group \"eval\"\n [\n test \"at 0.0\" test_eval_at_zero;\n test \"at 1.0\" test_eval_at_one;\n test \"negative clamped\" test_eval_negative_clamped;\n test \"above 1.0 clamped\" test_eval_above_one_clamped;\n ];\n group \"of_colors\"\n [\n test \"two stops midpoint\" test_two_stops_midpoint;\n test \"three stops midpoint\" test_three_stops_midpoint;\n test \"one stop raises\" test_one_stop_raises;\n test \"empty raises\" test_empty_raises;\n ];\n group \"predefined\"\n [\n test \"all evaluate without error\" test_predefined_no_raise;\n test \"viridis endpoints\" test_viridis_endpoints;\n ];\n ]\n"
},
{
"path": "packages/hugin/test/test_color.ml",
"content": "open Hugin\nopen Windtrap\n\nlet rgba_testable = quad (float 0.01) (float 0.01) (float 0.01) (float 0.01)\nlet check_float msg = equal ~msg (float 1e-6)\nlet check_rgba msg expected actual = equal ~msg rgba_testable expected actual\n\n(* sRGB roundtrip *)\n\nlet test_black_roundtrip () =\n let r, g, b, a = Color.rgb ~r:0. ~g:0. ~b:0. () |> Color.to_rgba in\n check_rgba \"black\" (0., 0., 0., 1.) (r, g, b, a)\n\nlet test_white_roundtrip () =\n let r, g, b, a = Color.rgb ~r:1. ~g:1. ~b:1. () |> Color.to_rgba in\n check_rgba \"white\" (1., 1., 1., 1.) (r, g, b, a)\n\nlet test_red_roundtrip () =\n let r, g, b, a = Color.rgb ~r:1. ~g:0. ~b:0. () |> Color.to_rgba in\n check_rgba \"red\" (1., 0., 0., 1.) (r, g, b, a)\n\nlet test_green_roundtrip () =\n let r, g, b, a = Color.rgb ~r:0. ~g:1. ~b:0. () |> Color.to_rgba in\n check_rgba \"green\" (0., 1., 0., 1.) (r, g, b, a)\n\nlet test_blue_roundtrip () =\n let r, g, b, a = Color.rgb ~r:0. ~g:0. ~b:1. () |> Color.to_rgba in\n (* Pure blue is near the sRGB gamut boundary in OKLCH, so the roundtrip is not\n perfectly lossless due to gamut clamping. We verify it's close. *)\n let wide_rgba = quad (float 0.15) (float 0.01) (float 0.01) (float 0.01) in\n equal ~msg:\"blue\" wide_rgba (0., 0., 1., 1.) (r, g, b, a)\n\nlet test_mid_gray_roundtrip () =\n let r, g, b, a = Color.rgb ~r:0.5 ~g:0.5 ~b:0.5 () |> Color.to_rgba in\n check_rgba \"mid gray\" (0.5, 0.5, 0.5, 1.) (r, g, b, a)\n\nlet test_arbitrary_roundtrip () =\n let r, g, b, a = Color.rgb ~r:0.3 ~g:0.6 ~b:0.9 () |> Color.to_rgba in\n (* OKLCH roundtrip has small error due to gamut boundary clamping *)\n let wide_rgba = quad (float 0.03) (float 0.03) (float 0.03) (float 0.01) in\n equal ~msg:\"arbitrary\" wide_rgba (0.3, 0.6, 0.9, 1.) (r, g, b, a)\n\n(* Hex parsing *)\n\nlet test_hex_6_with_hash () =\n let r, g, b, _ = Color.hex \"#FF0000\" |> Color.to_rgba in\n check_rgba \"red hex\" (1., 0., 0., 1.) (r, g, b, 1.)\n\nlet test_hex_6_without_hash () =\n let r, g, b, _ = Color.hex \"FF0000\" |> Color.to_rgba in\n check_rgba \"red hex no hash\" (1., 0., 0., 1.) (r, g, b, 1.)\n\nlet test_hex_8_with_alpha () =\n let _, _, _, a = Color.hex \"#FF000080\" |> Color.to_rgba in\n let expected_a = 128. /. 255. in\n check_float \"alpha\" expected_a a\n\nlet test_hex_case_insensitive () =\n let r1, g1, b1, _ = Color.hex \"#ff0000\" |> Color.to_rgba in\n let r2, g2, b2, _ = Color.hex \"#FF0000\" |> Color.to_rgba in\n check_rgba \"case insensitive\" (r1, g1, b1, 1.) (r2, g2, b2, 1.)\n\nlet test_hex_invalid_length () =\n raises_invalid_arg \"Color.hex: expected 6 or 8 hex digits, got 3\" (fun () ->\n Color.hex \"#FFF\")\n\nlet test_hex_invalid_chars () =\n raises_invalid_arg \"Color.hex: invalid hex digit 'G'\" (fun () ->\n Color.hex \"#GGGGGG\")\n\nlet test_hex_empty () =\n raises_invalid_arg \"Color.hex: expected 6 or 8 hex digits, got 0\" (fun () ->\n Color.hex \"\")\n\n(* OKLCH constructors *)\n\nlet test_oklch_fields () =\n let c = Color.oklch ~l:0.5 ~c:0.1 ~h:180. () in\n check_float \"lightness\" 0.5 (Color.lightness c);\n check_float \"chroma\" 0.1 (Color.chroma c);\n check_float \"hue\" 180. (Color.hue c);\n check_float \"alpha\" 1. (Color.alpha c)\n\nlet test_oklcha_alpha () =\n let c = Color.oklcha ~l:0.5 ~c:0.1 ~h:180. ~a:0.7 () in\n check_float \"alpha\" 0.7 (Color.alpha c)\n\n(* Operations *)\n\nlet test_lighten_clamps () =\n let c = Color.lighten 2.0 (Color.oklch ~l:0.8 ~c:0. ~h:0. ()) in\n check_float \"clamped to 1\" 1.0 (Color.lightness c)\n\nlet test_darken_clamps () =\n let c = Color.darken 2.0 (Color.oklch ~l:0.2 ~c:0. ~h:0. ()) in\n check_float \"clamped to 0\" 0.0 (Color.lightness c)\n\nlet test_lighten_adds () =\n let c = Color.lighten 0.1 (Color.oklch ~l:0.5 ~c:0. ~h:0. ()) in\n check_float \"lighten adds\" 0.6 (Color.lightness c)\n\nlet test_darken_subtracts () =\n let c = Color.darken 0.1 (Color.oklch ~l:0.5 ~c:0. ~h:0. ()) in\n check_float \"darken subtracts\" 0.4 (Color.lightness c)\n\nlet test_with_alpha () =\n let c = Color.with_alpha 0.3 (Color.oklch ~l:0.5 ~c:0. ~h:0. ()) in\n check_float \"with_alpha\" 0.3 (Color.alpha c)\n\n(* Mix *)\n\nlet a = Color.oklch ~l:0.2 ~c:0.05 ~h:10. ()\nlet b = Color.oklch ~l:0.8 ~c:0.15 ~h:50. ()\n\nlet test_mix_zero () =\n let m = Color.mix 0.0 a b in\n check_float \"lightness\" (Color.lightness a) (Color.lightness m);\n check_float \"chroma\" (Color.chroma a) (Color.chroma m);\n check_float \"hue\" (Color.hue a) (Color.hue m)\n\nlet test_mix_one () =\n let m = Color.mix 1.0 a b in\n check_float \"lightness\" (Color.lightness b) (Color.lightness m);\n check_float \"chroma\" (Color.chroma b) (Color.chroma m);\n check_float \"hue\" (Color.hue b) (Color.hue m)\n\nlet test_mix_midpoint_lightness () =\n let m = Color.mix 0.5 a b in\n check_float \"midpoint lightness\" 0.5 (Color.lightness m)\n\nlet test_mix_midpoint_chroma () =\n let m = Color.mix 0.5 a b in\n check_float \"midpoint chroma\" 0.1 (Color.chroma m)\n\nlet test_mix_hue_forward () =\n let c1 = Color.oklch ~l:0.5 ~c:0.1 ~h:10. () in\n let c2 = Color.oklch ~l:0.5 ~c:0.1 ~h:50. () in\n let m = Color.mix 0.5 c1 c2 in\n check_float \"hue forward\" 30. (Color.hue m)\n\nlet test_mix_hue_wraps_360 () =\n let c1 = Color.oklch ~l:0.5 ~c:0.1 ~h:350. () in\n let c2 = Color.oklch ~l:0.5 ~c:0.1 ~h:10. () in\n let m = Color.mix 0.5 c1 c2 in\n check_float \"hue wraps\" 0. (Color.hue m)\n\nlet test_mix_hue_reverse_wrap () =\n let c1 = Color.oklch ~l:0.5 ~c:0.1 ~h:10. () in\n let c2 = Color.oklch ~l:0.5 ~c:0.1 ~h:350. () in\n let m = Color.mix 0.5 c1 c2 in\n check_float \"hue reverse wrap\" 0. (Color.hue m)\n\nlet test_mix_alpha () =\n let c1 = Color.oklcha ~l:0.5 ~c:0. ~h:0. ~a:0.0 () in\n let c2 = Color.oklcha ~l:0.5 ~c:0. ~h:0. ~a:1.0 () in\n let m = Color.mix 0.5 c1 c2 in\n check_float \"alpha interpolates\" 0.5 (Color.alpha m)\n\n(* Gamut clamping *)\n\nlet test_high_chroma_clamped () =\n let r, g, b, _ = Color.oklch ~l:0.5 ~c:0.4 ~h:0. () |> Color.to_rgba in\n is_true ~msg:\"r in [0,1]\" (r >= 0. && r <= 1.);\n is_true ~msg:\"g in [0,1]\" (g >= 0. && g <= 1.);\n is_true ~msg:\"b in [0,1]\" (b >= 0. && b <= 1.)\n\n(* Named colors *)\n\nlet test_black_named () =\n let r, g, b, a = Color.black |> Color.to_rgba in\n check_rgba \"black\" (0., 0., 0., 1.) (r, g, b, a)\n\nlet test_white_named () =\n let r, g, b, a = Color.white |> Color.to_rgba in\n check_rgba \"white\" (1., 1., 1., 1.) (r, g, b, a)\n\nlet test_orange_matches_hex () =\n let r1, g1, b1, _ = Color.orange |> Color.to_rgba in\n let r2, g2, b2, _ = Color.hex \"#E69F00\" |> Color.to_rgba in\n check_rgba \"orange\" (r1, g1, b1, 1.) (r2, g2, b2, 1.)\n\nlet () =\n run \"Color\"\n [\n group \"sRGB roundtrip\"\n [\n test \"black\" test_black_roundtrip;\n test \"white\" test_white_roundtrip;\n test \"red\" test_red_roundtrip;\n test \"green\" test_green_roundtrip;\n test \"blue\" test_blue_roundtrip;\n test \"mid gray\" test_mid_gray_roundtrip;\n test \"arbitrary\" test_arbitrary_roundtrip;\n ];\n group \"hex parsing\"\n [\n test \"6-digit with hash\" test_hex_6_with_hash;\n test \"6-digit without hash\" test_hex_6_without_hash;\n test \"8-digit with alpha\" test_hex_8_with_alpha;\n test \"case insensitive\" test_hex_case_insensitive;\n test \"invalid length\" test_hex_invalid_length;\n test \"invalid chars\" test_hex_invalid_chars;\n test \"empty string\" test_hex_empty;\n ];\n group \"OKLCH constructors\"\n [\n test \"oklch fields\" test_oklch_fields;\n test \"oklcha alpha\" test_oklcha_alpha;\n ];\n group \"operations\"\n [\n test \"lighten clamps\" test_lighten_clamps;\n test \"darken clamps\" test_darken_clamps;\n test \"lighten adds\" test_lighten_adds;\n test \"darken subtracts\" test_darken_subtracts;\n test \"with_alpha\" test_with_alpha;\n ];\n group \"mix\"\n [\n test \"mix 0.0 returns first\" test_mix_zero;\n test \"mix 1.0 returns second\" test_mix_one;\n test \"midpoint lightness\" test_mix_midpoint_lightness;\n test \"midpoint chroma\" test_mix_midpoint_chroma;\n test \"hue shortest arc forward\" test_mix_hue_forward;\n test \"hue wraps across 360\" test_mix_hue_wraps_360;\n test \"hue reverse wrap\" test_mix_hue_reverse_wrap;\n test \"alpha interpolates\" test_mix_alpha;\n ];\n group \"gamut clamping\"\n [ test \"high chroma clamped\" test_high_chroma_clamped ];\n group \"named colors\"\n [\n test \"black\" test_black_named;\n test \"white\" test_white_named;\n test \"orange matches hex\" test_orange_matches_hex;\n ];\n ]\n"
},
{
"path": "packages/hugin/test/test_image_util.ml",
"content": "(*---------------------------------------------------------------------------\n Tests for Image_util.base64_encode — exercised indirectly through Hugin.pp\n which calls base64_encode on PNG buffer data. We also test the base64 logic\n through the pp data URI output.\n ---------------------------------------------------------------------------*)\n\nopen Hugin\nopen Windtrap\n\nlet contains s sub =\n let len_s = String.length s and len_sub = String.length sub in\n if len_sub > len_s then false\n else\n let found = ref false in\n for i = 0 to len_s - len_sub do\n if (not !found) && String.sub s i len_sub = sub then found := true\n done;\n !found\n\nlet is_base64_char = function\n | 'A' .. 'Z' | 'a' .. 'z' | '0' .. '9' | '+' | '/' | '=' -> true\n | _ -> false\n\nlet sample_x = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0))\nlet sample_y = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0))\n\n(* pp produces a data URI with base64 encoded PNG *)\n\nlet test_pp_data_uri () =\n let spec = Hugin.line ~x:sample_x ~y:sample_y () in\n let buf = Buffer.create 256 in\n let fmt = Format.formatter_of_buffer buf in\n Hugin.pp fmt spec;\n Format.pp_print_flush fmt ();\n let output = Buffer.contents buf in\n is_true ~msg:\"starts with image markdown\"\n (contains output \"![figure](data:image/png;base64,\");\n (* Base64 output should only contain valid base64 chars *)\n let b64_start = \"base64,\" in\n let start_idx =\n let rec find i =\n if i > String.length output - String.length b64_start then -1\n else if String.sub output i (String.length b64_start) = b64_start then\n i + String.length b64_start\n else find (i + 1)\n in\n find 0\n in\n is_true ~msg:\"found base64 data\" (start_idx > 0);\n if start_idx > 0 then begin\n let end_idx = String.length output - 1 in\n let b64 = String.sub output start_idx (end_idx - start_idx) in\n is_true ~msg:\"all chars are valid base64\"\n (String.to_seq b64 |> Seq.for_all is_base64_char)\n end\n\n(* render_to_buffer produces non-empty data *)\n\nlet test_render_to_buffer () =\n let spec = Hugin.line ~x:sample_x ~y:sample_y () in\n let buf = Hugin.render_to_buffer spec in\n is_true ~msg:\"non-empty\" (String.length buf > 0);\n (* PNG magic bytes: 0x89 P N G *)\n is_true ~msg:\"PNG magic byte\" (Char.code (String.get buf 0) = 0x89);\n is_true ~msg:\"PNG P\" (String.get buf 1 = 'P');\n is_true ~msg:\"PNG N\" (String.get buf 2 = 'N');\n is_true ~msg:\"PNG G\" (String.get buf 3 = 'G')\n\nlet () =\n run \"Image_util\"\n [\n group \"pp data URI\"\n [ test \"produces valid base64 data URI\" test_pp_data_uri ];\n group \"render_to_buffer\"\n [ test \"produces valid PNG\" test_render_to_buffer ];\n ]\n"
},
{
"path": "packages/hugin/test/test_resolve.ml",
"content": "open Hugin\nopen Windtrap\n\n(* Helpers *)\n\nlet contains s sub =\n let len_s = String.length s and len_sub = String.length sub in\n if len_sub > len_s then false\n else\n let found = ref false in\n for i = 0 to len_s - len_sub do\n if (not !found) && String.sub s i len_sub = sub then found := true\n done;\n !found\n\nlet count_substring s sub =\n let len_s = String.length s and len_sub = String.length sub in\n if len_sub > len_s || len_sub = 0 then 0\n else begin\n let count = ref 0 in\n for i = 0 to len_s - len_sub do\n if String.sub s i len_sub = sub then incr count\n done;\n !count\n end\n\nlet render ?(width = 400.) ?(height = 300.) spec =\n let tmp = Filename.temp_file \"hugin_test\" \".svg\" in\n Hugin.render_svg ~width ~height tmp spec;\n let ic = open_in tmp in\n let n = in_channel_length ic in\n let s = really_input_string ic n in\n close_in ic;\n Sys.remove tmp;\n s\n\nlet sample_x = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0))\nlet sample_y = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0) *. 2.)\nlet sample_line () = Hugin.line ~x:sample_x ~y:sample_y ()\nlet sample_point () = Hugin.point ~x:sample_x ~y:sample_y ()\n\nlet sample_bar () =\n Hugin.bar ~x:sample_x\n ~height:(Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0) +. 1.))\n ()\n\n(* basic marks *)\n\nlet test_line_resolves () =\n let svg = render (sample_line ()) in\n is_true ~msg:\"starts with xml\"\n (String.length svg > 5 && String.sub svg 0 5 = \" float_of_int i.(0) /. 10.) in\n let svg = render (Hugin.hist ~x:data ()) in\n is_true ~msg:\"contains path\" (contains svg \"hello<\")\n\nlet test_hline_resolves () =\n let spec = Hugin.layers [ sample_line (); Hugin.hline ~y:3. () ] in\n let svg = render spec in\n (* hline adds a horizontal path; 2 paths from line+hline vs 1 without *)\n let path_count = count_substring svg \"= 2)\n\nlet test_vline_resolves () =\n let spec = Hugin.layers [ sample_line (); Hugin.vline ~x:2. () ] in\n let svg = render spec in\n let path_count = count_substring svg \"= 2)\n\nlet test_empty_layers () =\n let svg = render (Hugin.layers []) in\n is_true ~msg:\"valid svg\" (contains svg \" Hugin.title \"My Title\") in\n is_true ~msg:\"title in svg\" (contains svg \">My Title<\")\n\nlet test_xlabel_appears () =\n let svg = render (sample_line () |> Hugin.xlabel \"X Axis\") in\n is_true ~msg:\"xlabel in svg\" (contains svg \">X Axis<\")\n\nlet test_ylabel_appears () =\n let svg = render (sample_line () |> Hugin.ylabel \"Y Axis\") in\n is_true ~msg:\"ylabel in svg\" (contains svg \">Y Axis<\")\n\nlet test_outermost_title_wins () =\n (* decorate prepends to the decoration list, and apply_decoration keeps the\n first-seen title. So the outermost (last-applied) title wins. *)\n let svg =\n render (sample_line () |> Hugin.title \"Inner\" |> Hugin.title \"Outer\")\n in\n is_true ~msg:\"outer title present\" (contains svg \">Outer<\");\n is_false ~msg:\"inner title absent\" (contains svg \">Inner<\")\n\n(* histogram normalization *)\n\nlet test_hist_bins () =\n let data = Nx.init Float32 [| 100 |] (fun i -> float_of_int i.(0)) in\n let svg = render (Hugin.hist ~x:data ~bins:(`Num 5) ()) in\n is_true ~msg:\"produces paths\" (contains svg \" float_of_int i.(0)) in\n let svg = render (Hugin.hist ~x:data ~bins:(`Num 5) ~density:true ()) in\n (* density normalization should produce bars; 5 bins = 5 closed paths *)\n let z_count = count_substring svg \" Z\\\"\" in\n is_true ~msg:\"has 5 bars\" (z_count >= 5)\n\nlet test_hist_edges () =\n let data = Nx.init Float32 [| 100 |] (fun i -> float_of_int i.(0)) in\n let svg = render (Hugin.hist ~x:data ~bins:(`Edges [| 0.; 50.; 100. |]) ()) in\n (* 2 bins from 3 edges = 2 closed paths *)\n let z_count = count_substring svg \" Z\\\"\" in\n is_true ~msg:\"has 2 bars\" (z_count >= 2)\n\n(* auto coloring *)\n\nlet test_auto_color_different () =\n let line1 = Hugin.line ~x:sample_x ~y:sample_y () in\n let y2 = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0) *. 3.) in\n let line2 = Hugin.line ~x:sample_x ~y:y2 () in\n let svg = render (Hugin.layers [ line1; line2 ]) in\n let stroke_count = count_substring svg \"stroke=\\\"rgb(\" in\n is_true ~msg:\"multiple stroke colors\" (stroke_count >= 2)\n\nlet test_explicit_color_preserved () =\n let svg = render (Hugin.line ~x:sample_x ~y:sample_y ~color:Color.black ()) in\n is_true ~msg:\"has black stroke\" (contains svg \"stroke=\\\"rgb(0,0,0)\\\"\")\n\n(* grid layout *)\n\nlet test_grid_2x2 () =\n let a = sample_line () |> Hugin.title \"A\" in\n let b = sample_line () |> Hugin.title \"B\" in\n let c = sample_line () |> Hugin.title \"C\" in\n let d = sample_line () |> Hugin.title \"D\" in\n let svg = render (Hugin.grid [ [ a; b ]; [ c; d ] ]) in\n is_true ~msg:\"has A\" (contains svg \">A<\");\n is_true ~msg:\"has D\" (contains svg \">D<\");\n (* 4 panels = 4 clip regions *)\n let clip_count = count_substring svg \" Hugin.title \"L\" in\n let b = sample_line () |> Hugin.title \"R\" in\n let svg = render (Hugin.hstack [ a; b ]) in\n is_true ~msg:\"has L\" (contains svg \">L<\");\n is_true ~msg:\"has R\" (contains svg \">R<\")\n\nlet test_vstack () =\n let a = sample_line () |> Hugin.title \"Top\" in\n let b = sample_line () |> Hugin.title \"Bot\" in\n let svg = render (Hugin.vstack [ a; b ]) in\n is_true ~msg:\"has Top\" (contains svg \">Top<\");\n is_true ~msg:\"has Bot\" (contains svg \">Bot<\")\n\n(* themes *)\n\nlet test_dark_theme () =\n let svg = render (sample_line () |> Hugin.with_theme Theme.dark) in\n (* dark theme has dark background — rgb values near 0 *)\n is_true ~msg:\"has dark fill\" (contains svg \"fill=\\\"rgb(\");\n is_true ~msg:\"has light strokes\" (contains svg \"stroke=\\\"rgb(\")\n\nlet test_minimal_theme () =\n let svg_default = render (sample_line ()) in\n let svg_minimal = render (sample_line () |> Hugin.with_theme Theme.minimal) in\n (* minimal theme has no grid, so fewer paths *)\n let default_paths = count_substring svg_default \" Hugin.grid_lines true) in\n let svg_off = render (sample_line () |> Hugin.grid_lines false) in\n let on_paths = count_substring svg_on \" Hugin.legend) in\n is_true ~msg:\"legend text\" (contains svg \">Series A<\")\n\n(* fill_between *)\n\nlet test_fill_between_resolves () =\n let y2 = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0) *. 3.) in\n let svg = render (Hugin.fill_between ~x:sample_x ~y1:sample_y ~y2 ()) in\n is_true ~msg:\"contains path\" (contains svg \" float_of_int i.(0) *. 3.) in\n let spec =\n Hugin.fill_between ~x:sample_x ~y1:sample_y ~y2 ~label:\"band\" ()\n |> Hugin.legend\n in\n let svg = render spec in\n is_true ~msg:\"legend text\" (contains svg \">band<\")\n\n(* hspan / vspan *)\n\nlet test_hspan_resolves () =\n let svg_base = render (sample_line ()) in\n let spec = Hugin.layers [ sample_line (); Hugin.hspan ~y0:1. ~y1:3. () ] in\n let svg = render spec in\n (* hspan adds a filled rectangle = one extra closed path *)\n let base_z = count_substring svg_base \" Z\\\"\" in\n let with_z = count_substring svg \" Z\\\"\" in\n is_true ~msg:\"more closed paths with hspan\" (with_z > base_z)\n\nlet test_vspan_resolves () =\n let svg_base = render (sample_line ()) in\n let spec = Hugin.layers [ sample_line (); Hugin.vspan ~x0:1. ~x1:3. () ] in\n let svg = render spec in\n let base_z = count_substring svg_base \" Z\\\"\" in\n let with_z = count_substring svg \" Z\\\"\" in\n is_true ~msg:\"more closed paths with vspan\" (with_z > base_z)\n\n(* step line *)\n\nlet test_step_post () =\n let svg_normal = render (Hugin.line ~x:sample_x ~y:sample_y ()) in\n let svg_step = render (Hugin.line ~x:sample_x ~y:sample_y ~step:`Post ()) in\n (* step line inserts intermediate points, so more L commands in total *)\n let normal_l = count_substring svg_normal \" L\" in\n let step_l = count_substring svg_step \" L\" in\n is_true ~msg:\"step has more L commands\" (step_l > normal_l)\n\nlet test_step_pre () =\n let svg_normal = render (Hugin.line ~x:sample_x ~y:sample_y ()) in\n let svg_step = render (Hugin.line ~x:sample_x ~y:sample_y ~step:`Pre ()) in\n (* pre step also inserts intermediate points *)\n let normal_l = count_substring svg_normal \" L\" in\n let step_l = count_substring svg_step \" L\" in\n is_true ~msg:\"step has more L commands\" (step_l > normal_l)\n\nlet test_step_mid () =\n let svg_normal = render (Hugin.line ~x:sample_x ~y:sample_y ()) in\n let svg_step = render (Hugin.line ~x:sample_x ~y:sample_y ~step:`Mid ()) in\n (* mid step inserts 2 intermediate points per segment *)\n let normal_l = count_substring svg_normal \" L\" in\n let step_l = count_substring svg_step \" L\" in\n is_true ~msg:\"step has more L commands\" (step_l > normal_l)\n\n(* errorbar *)\n\nlet test_errorbar_symmetric () =\n let err = Nx.init Float32 [| 5 |] (fun _ -> 0.5) in\n let svg =\n render (Hugin.errorbar ~x:sample_x ~y:sample_y ~yerr:(`Symmetric err) ())\n in\n (* 5 points × 3 paths each (stem + 2 caps) = 15 paths *)\n let path_count = count_substring svg \"= 15)\n\nlet test_errorbar_asymmetric () =\n let elo = Nx.init Float32 [| 5 |] (fun _ -> 0.3) in\n let ehi = Nx.init Float32 [| 5 |] (fun _ -> 0.7) in\n let svg =\n render\n (Hugin.errorbar ~x:sample_x ~y:sample_y ~yerr:(`Asymmetric (elo, ehi)) ())\n in\n let path_count = count_substring svg \"= 15)\n\nlet test_errorbar_with_xerr () =\n let yerr = Nx.init Float32 [| 5 |] (fun _ -> 0.5) in\n let xerr = Nx.init Float32 [| 5 |] (fun _ -> 0.2) in\n let svg =\n render\n (Hugin.errorbar ~x:sample_x ~y:sample_y ~yerr:(`Symmetric yerr)\n ~xerr:(`Symmetric xerr) ())\n in\n (* with xerr: 5 points × 6 paths each (yerr stem+2caps + xerr stem+2caps) =\n 30 *)\n let svg_yerr_only =\n render (Hugin.errorbar ~x:sample_x ~y:sample_y ~yerr:(`Symmetric yerr) ())\n in\n let yerr_paths = count_substring svg_yerr_only \" yerr_paths)\n\n(* heatmap *)\n\nlet test_heatmap_resolves () =\n let data =\n Nx.init Float32 [| 3; 4 |] (fun i -> float_of_int (i.(0) + i.(1)))\n in\n let svg = render (Hugin.heatmap ~data ()) in\n is_true ~msg:\"contains paths\" (contains svg \" float_of_int (i.(0) + i.(1)))\n in\n let svg = render (Hugin.heatmap ~data ~annotate:true ()) in\n is_true ~msg:\"contains text\" (contains svg \" 0.5) in\n let svg =\n render\n (Hugin.heatmap ~data ~annotate:true\n ~fmt:(fun v -> Printf.sprintf \"%.0f%%\" (v *. 100.))\n ())\n in\n is_true ~msg:\"contains formatted text\" (contains svg \">50%<\")\n\n(* imshow *)\n\nlet imshow_data () =\n Nx.init Float32 [| 4; 6 |] (fun i -> float_of_int i.(0) +. float_of_int i.(1))\n\nlet test_imshow_rasterizes_to_image () =\n let svg = render (Hugin.imshow ~data:(imshow_data ()) ()) in\n (* imshow is rasterized to an Image in the Prepared stage — verify the SVG\n backend emits an element with base64 PNG data *)\n is_true ~msg:\"contains image element\" (contains svg \" svg_linear);\n is_true ~msg:\"sqrt differs from linear\" (svg_sqrt <> svg_linear);\n is_true ~msg:\"log differs from sqrt\" (svg_log <> svg_sqrt)\n\nlet test_imshow_cmap_changes_output () =\n let data = imshow_data () in\n let svg_default = render (Hugin.imshow ~data ()) in\n let svg_hot = render (Hugin.imshow ~data ~cmap:Cmap.hot ()) in\n let svg_gray = render (Hugin.imshow ~data ~cmap:Cmap.gray_r ()) in\n is_true ~msg:\"hot differs from default\" (svg_hot <> svg_default);\n is_true ~msg:\"gray_r differs from hot\" (svg_gray <> svg_hot)\n\n(* contour *)\n\nlet contour_data () =\n (* Concentric circles centered at (4.5, 4.5), values = r² *)\n Nx.init Float32 [| 10; 10 |] (fun i ->\n let x = float_of_int i.(1) -. 4.5 in\n let y = float_of_int i.(0) -. 4.5 in\n (x *. x) +. (y *. y))\n\nlet test_contour_unfilled_has_stroked_paths () =\n let svg =\n render\n (Hugin.contour ~data:(contour_data ()) ~x0:0. ~x1:9. ~y0:0. ~y1:9.\n ~levels:(`Num 4) ())\n in\n (* Unfilled contours are stroked paths (stroke=, fill=\"none\") *)\n is_true ~msg:\"has stroked paths\" (contains svg \"stroke=\\\"rgb(\");\n let path_count = count_substring svg \"= 4)\n\nlet test_contour_filled_more_paths () =\n let data = contour_data () in\n let svg_unfilled =\n render (Hugin.contour ~data ~x0:0. ~x1:9. ~y0:0. ~y1:9. ~levels:(`Num 4) ())\n in\n let svg_filled =\n render\n (Hugin.contour ~data ~x0:0. ~x1:9. ~y0:0. ~y1:9. ~levels:(`Num 4)\n ~filled:true ())\n in\n (* Filled contours use fill=rgb(...), unfilled use stroke *)\n is_true ~msg:\"filled has fill color\" (contains svg_filled \"fill=\\\"rgb(\");\n (* Filled output differs from unfilled *)\n is_true ~msg:\"filled differs from unfilled\" (svg_filled <> svg_unfilled)\n\nlet test_contour_level_count_affects_paths () =\n let data = contour_data () in\n let svg_few =\n render (Hugin.contour ~data ~x0:0. ~x1:9. ~y0:0. ~y1:9. ~levels:(`Num 2) ())\n in\n let svg_many =\n render (Hugin.contour ~data ~x0:0. ~x1:9. ~y0:0. ~y1:9. ~levels:(`Num 8) ())\n in\n let few_paths = count_substring svg_few \" few_paths)\n\nlet test_contour_legend () =\n let svg =\n render\n (Hugin.contour ~data:(contour_data ()) ~x0:0. ~x1:9. ~y0:0. ~y1:9.\n ~label:\"density\" ()\n |> Hugin.legend)\n in\n is_true ~msg:\"legend text\" (contains svg \">density<\")\n\n(* inverted axes *)\n\nlet test_invert_changes_path_data () =\n (* Inversion reverses the scale mapping, so the path d= attribute must differ\n between normal and inverted rendering of the same data. *)\n let svg_normal = render (sample_line ()) in\n let svg_xinv = render (sample_line () |> Hugin.xinvert) in\n let svg_yinv = render (sample_line () |> Hugin.yinvert) in\n is_true ~msg:\"xinvert changes path\" (svg_xinv <> svg_normal);\n is_true ~msg:\"yinvert changes path\" (svg_yinv <> svg_normal)\n\nlet test_yinvert_hr_diagram () =\n (* An HR diagram uses yinvert (brighter stars at top) and decorations. *)\n let bv = Nx.create Float32 [| 5 |] [| -0.3; 0.; 0.5; 1.0; 1.5 |] in\n let mag = Nx.create Float32 [| 5 |] [| -5.; 0.; 2.; 5.; 10. |] in\n let svg =\n render\n (Hugin.point ~x:bv ~y:mag ()\n |> Hugin.yinvert |> Hugin.xlabel \"B-V\" |> Hugin.ylabel \"Magnitude\")\n in\n is_true ~msg:\"xlabel present\" (contains svg \">B-V<\");\n is_true ~msg:\"ylabel present\" (contains svg \">Magnitude<\");\n is_true ~msg:\"has markers\" (contains svg \" Hugin.xtick_format (fun v -> Printf.sprintf \"%.0f%%\" (v *. 100.))\n in\n let svg = render spec in\n (* x data is 0..4, so formatted ticks should contain \"%\" *)\n is_true ~msg:\"formatted ticks contain %\" (contains svg \"%\")\n\nlet test_ytick_format () =\n let spec =\n sample_line () |> Hugin.ytick_format (fun v -> Printf.sprintf \"$%.0f\" v)\n in\n let svg = render spec in\n (* y data is 0..8, so formatted ticks should contain \"$\" *)\n is_true ~msg:\"formatted ticks contain $\" (contains svg \"$\")\n\nlet () =\n run \"Resolve\"\n [\n group \"basic marks\"\n [\n test \"line\" test_line_resolves;\n test \"point\" test_point_resolves;\n test \"bar\" test_bar_resolves;\n test \"hist\" test_hist_resolves;\n test \"text\" test_text_mark_resolves;\n test \"hline\" test_hline_resolves;\n test \"vline\" test_vline_resolves;\n test \"empty layers\" test_empty_layers;\n ];\n group \"decorations\"\n [\n test \"title appears\" test_title_appears;\n test \"xlabel appears\" test_xlabel_appears;\n test \"ylabel appears\" test_ylabel_appears;\n test \"outermost title wins\" test_outermost_title_wins;\n ];\n group \"histogram normalization\"\n [\n test \"bins\" test_hist_bins;\n test \"density\" test_hist_density;\n test \"edges\" test_hist_edges;\n ];\n group \"auto coloring\"\n [\n test \"different colors\" test_auto_color_different;\n test \"explicit color preserved\" test_explicit_color_preserved;\n ];\n group \"grid layout\"\n [\n test \"2x2 grid\" test_grid_2x2;\n test \"empty grid\" test_grid_empty;\n test \"hstack\" test_hstack;\n test \"vstack\" test_vstack;\n ];\n group \"themes\"\n [\n test \"dark theme\" test_dark_theme;\n test \"minimal theme\" test_minimal_theme;\n ];\n group \"grid lines\" [ test \"grid lines off\" test_grid_lines_off ];\n group \"legend\" [ test \"legend appears\" test_legend_appears ];\n group \"fill_between\"\n [\n test \"resolves\" test_fill_between_resolves;\n test \"with label\" test_fill_between_with_label;\n ];\n group \"hspan/vspan\"\n [ test \"hspan\" test_hspan_resolves; test \"vspan\" test_vspan_resolves ];\n group \"step line\"\n [\n test \"post\" test_step_post;\n test \"pre\" test_step_pre;\n test \"mid\" test_step_mid;\n ];\n group \"errorbar\"\n [\n test \"symmetric\" test_errorbar_symmetric;\n test \"asymmetric\" test_errorbar_asymmetric;\n test \"with xerr\" test_errorbar_with_xerr;\n ];\n group \"heatmap\"\n [\n test \"resolves\" test_heatmap_resolves;\n test \"annotated\" test_heatmap_annotated;\n test \"custom fmt\" test_heatmap_custom_fmt;\n ];\n group \"tick format\"\n [\n test \"xtick_format\" test_xtick_format;\n test \"ytick_format\" test_ytick_format;\n ];\n group \"imshow\"\n [\n test \"rasterizes to image\" test_imshow_rasterizes_to_image;\n test \"stretches differ\" test_imshow_stretches_differ;\n test \"cmap changes output\" test_imshow_cmap_changes_output;\n ];\n group \"contour\"\n [\n test \"unfilled has stroked paths\"\n test_contour_unfilled_has_stroked_paths;\n test \"filled more paths\" test_contour_filled_more_paths;\n test \"level count affects paths\"\n test_contour_level_count_affects_paths;\n test \"legend\" test_contour_legend;\n ];\n group \"inverted axes\"\n [\n test \"invert changes path data\" test_invert_changes_path_data;\n test \"yinvert HR diagram\" test_yinvert_hr_diagram;\n ];\n ]\n"
},
{
"path": "packages/hugin/test/test_scale.ml",
"content": "(*---------------------------------------------------------------------------\n Tests for Scale logic — exercised indirectly through Hugin.render_svg. We\n verify that linear and log scales produce correct axis tick labels in the SVG\n output, which proves the scale math is correct.\n ---------------------------------------------------------------------------*)\n\nopen Hugin\nopen Windtrap\n\nlet contains s sub =\n let len_s = String.length s and len_sub = String.length sub in\n if len_sub > len_s then false\n else\n let found = ref false in\n for i = 0 to len_s - len_sub do\n if (not !found) && String.sub s i len_sub = sub then found := true\n done;\n !found\n\nlet count_substring s sub =\n let len_s = String.length s and len_sub = String.length sub in\n if len_sub > len_s || len_sub = 0 then 0\n else begin\n let count = ref 0 in\n for i = 0 to len_s - len_sub do\n if String.sub s i len_sub = sub then incr count\n done;\n !count\n end\n\nlet render spec =\n let tmp = Filename.temp_file \"hugin_test\" \".svg\" in\n Hugin.render_svg ~width:400. ~height:300. tmp spec;\n let ic = open_in tmp in\n let n = in_channel_length ic in\n let s = really_input_string ic n in\n close_in ic;\n Sys.remove tmp;\n s\n\nlet x5 = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0))\nlet y5 = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0))\n\n(* linear scale *)\n\nlet test_linear_ticks_present () =\n let svg = render (Hugin.line ~x:x5 ~y:y5 ()) in\n (* Data range 0-4, auto-ticks should include 0 *)\n is_true ~msg:\"has tick 0\" (contains svg \">0<\")\n\nlet test_linear_xlim () =\n (* Use different x and y ranges so we can distinguish x ticks from y ticks. x\n data: 0..10, y data: 100..200. With xlim 0-5, x ticks stay in [0,5] but y\n ticks are around 100-200 — no overlap. *)\n let x = Nx.init Float32 [| 11 |] (fun i -> float_of_int i.(0)) in\n let y =\n Nx.init Float32 [| 11 |] (fun i -> 100. +. (float_of_int i.(0) *. 10.))\n in\n let svg = render (Hugin.line ~x ~y () |> Hugin.xlim 0. 5.) in\n is_true ~msg:\"has tick 0\" (contains svg \">0<\");\n (* With xlim 0-5, we should not see x-axis tick \"8\" or \"10\". Y-axis ticks are\n in 100-200 range so no confusion. *)\n is_false ~msg:\"no tick 8\" (contains svg \">8<\");\n is_false ~msg:\"no tick 10\" (contains svg \">10<\")\n\nlet test_linear_ylim () =\n let x = Nx.init Float32 [| 11 |] (fun i -> float_of_int i.(0)) in\n let y = Nx.init Float32 [| 11 |] (fun i -> float_of_int i.(0)) in\n let svg = render (Hugin.line ~x ~y () |> Hugin.ylim 0. 5.) in\n is_true ~msg:\"valid svg\" (contains svg \"0<\")\n\nlet test_linear_small_range () =\n let x = Nx.create Float32 [| 3 |] [| 0.; 0.0005; 0.001 |] in\n let y = Nx.create Float32 [| 3 |] [| 0.; 0.5; 1. |] in\n let svg = render (Hugin.line ~x ~y ()) in\n is_true ~msg:\"valid svg\" (contains svg \" Hugin.xscale `Log) in\n is_true ~msg:\"has tick 1\" (contains svg \">1<\");\n is_true ~msg:\"has tick 10\" (contains svg \">10<\");\n is_true ~msg:\"has tick 10^2\" (contains svg \">10^2<\");\n is_true ~msg:\"has tick 10^3\" (contains svg \">10^3<\")\n\nlet test_log_y () =\n let x = Nx.create Float32 [| 4 |] [| 1.; 2.; 3.; 4. |] in\n let y = Nx.create Float32 [| 4 |] [| 1.; 10.; 100.; 1000. |] in\n let svg = render (Hugin.line ~x ~y () |> Hugin.yscale `Log) in\n is_true ~msg:\"has tick 1\" (contains svg \">1<\");\n is_true ~msg:\"has tick 10^3\" (contains svg \">10^3<\")\n\n(* custom ticks *)\n\nlet test_explicit_xticks () =\n let svg =\n render\n (Hugin.line ~x:x5 ~y:y5 () |> Hugin.xticks [ (0., \"zero\"); (4., \"four\") ])\n in\n is_true ~msg:\"has custom tick zero\" (contains svg \">zero<\");\n is_true ~msg:\"has custom tick four\" (contains svg \">four<\")\n\nlet test_explicit_yticks () =\n let svg =\n render\n (Hugin.line ~x:x5 ~y:y5 () |> Hugin.yticks [ (0., \"low\"); (4., \"high\") ])\n in\n is_true ~msg:\"has custom tick low\" (contains svg \">low<\");\n is_true ~msg:\"has custom tick high\" (contains svg \">high<\")\n\n(* sqrt scale *)\n\nlet test_sqrt_handles_zero () =\n (* Sqrt scale handles zero gracefully — critical for astronomical fluxes *)\n let x = Nx.create Float32 [| 5 |] [| 0.; 1.; 4.; 9.; 16. |] in\n let y = Nx.create Float32 [| 5 |] [| 0.; 1.; 2.; 3.; 4. |] in\n let svg = render (Hugin.line ~x ~y () |> Hugin.xscale `Sqrt) in\n is_true ~msg:\"has tick 0\" (contains svg \">0<\");\n is_true ~msg:\"has path\" (contains svg \" Hugin.yscale `Sqrt) in\n is_true ~msg:\"sqrt changes output\" (svg_sqrt <> svg_lin)\n\n(* asinh scale *)\n\nlet test_asinh_negative_values () =\n (* Asinh handles negative values, unlike log — needed for\n background-subtracted fluxes *)\n let x = Nx.create Float32 [| 5 |] [| -100.; -1.; 0.; 1.; 100. |] in\n let y = Nx.create Float32 [| 5 |] [| 0.; 1.; 2.; 3.; 4. |] in\n let svg = render (Hugin.line ~x ~y () |> Hugin.xscale `Asinh) in\n is_true ~msg:\"has tick 0\" (contains svg \">0<\");\n is_true ~msg:\"has path\" (contains svg \" Hugin.yscale `Asinh) in\n is_true ~msg:\"asinh changes output\" (svg_asinh <> svg_lin)\n\n(* symlog scale *)\n\nlet test_symlog_has_linear_and_log_ticks () =\n (* Symlog should produce ticks in both the linear region (near 0) and the log\n region (far from 0) *)\n let x =\n Nx.create Float32 [| 7 |] [| -1000.; -10.; -1.; 0.; 1.; 10.; 1000. |]\n in\n let y = Nx.init Float32 [| 7 |] (fun i -> float_of_int i.(0)) in\n let svg = render (Hugin.line ~x ~y () |> Hugin.xscale (`Symlog 10.)) in\n is_true ~msg:\"has tick 0 (linear region)\" (contains svg \">0<\");\n is_true ~msg:\"has path\" (contains svg \" float_of_int i.(0)) in\n let svg_lin = render (Hugin.line ~x ~y ()) in\n let svg_sym = render (Hugin.line ~x ~y () |> Hugin.xscale (`Symlog 10.)) in\n is_true ~msg:\"symlog changes output\" (svg_sym <> svg_lin)\n\n(* inverted scales *)\n\nlet test_invert_reverses_tick_order () =\n (* The same tick labels should appear, but xinvert swaps pixel positions. We\n verify the SVG output actually changes. *)\n let svg_normal = render (Hugin.line ~x:x5 ~y:y5 ()) in\n let svg_inv = render (Hugin.line ~x:x5 ~y:y5 () |> Hugin.xinvert) in\n is_true ~msg:\"has tick 0\" (contains svg_inv \">0<\");\n is_true ~msg:\"invert changes output\" (svg_inv <> svg_normal)\n\nlet test_invert_preserves_ticks () =\n (* Inversion should not remove or add ticks, just reposition them *)\n let svg_normal = render (Hugin.line ~x:x5 ~y:y5 ()) in\n let svg_inv = render (Hugin.line ~x:x5 ~y:y5 () |> Hugin.yinvert) in\n let normal_texts = count_substring svg_normal \" Hugin.xscale `Log |> Hugin.xinvert)\n in\n is_true ~msg:\"has tick 1\" (contains svg \">1<\");\n is_true ~msg:\"has tick 10\" (contains svg \">10<\")\n\nlet () =\n run \"Scale\"\n [\n group \"linear\"\n [\n test \"ticks present\" test_linear_ticks_present;\n test \"xlim constrains\" test_linear_xlim;\n test \"ylim constrains\" test_linear_ylim;\n test \"negative range\" test_linear_negative_range;\n test \"small range\" test_linear_small_range;\n test \"single point\" test_linear_single_point;\n ];\n group \"log\"\n [\n test \"power-of-10 ticks\" test_log_ticks; test \"log y axis\" test_log_y;\n ];\n group \"sqrt\"\n [\n test \"handles zero\" test_sqrt_handles_zero;\n test \"differs from linear\" test_sqrt_differs_from_linear;\n ];\n group \"asinh\"\n [\n test \"negative values\" test_asinh_negative_values;\n test \"differs from linear\" test_asinh_differs_from_linear;\n ];\n group \"symlog\"\n [\n test \"linear and log ticks\" test_symlog_has_linear_and_log_ticks;\n test \"differs from linear\" test_symlog_differs_from_linear;\n ];\n group \"inverted\"\n [\n test \"reverses tick order\" test_invert_reverses_tick_order;\n test \"preserves ticks\" test_invert_preserves_ticks;\n test \"log inverted\" test_log_inverted;\n ];\n group \"custom ticks\"\n [\n test \"explicit xticks\" test_explicit_xticks;\n test \"explicit yticks\" test_explicit_yticks;\n ];\n ]\n"
},
{
"path": "packages/hugin/test/test_svg_backend.ml",
"content": "(*---------------------------------------------------------------------------\n Tests for the SVG backend — rendered through Hugin.render_svg. We verify SVG\n structure, XML escaping, and content correctness.\n ---------------------------------------------------------------------------*)\n\nopen Hugin\nopen Windtrap\n\nlet contains s sub =\n let len_s = String.length s and len_sub = String.length sub in\n if len_sub > len_s then false\n else\n let found = ref false in\n for i = 0 to len_s - len_sub do\n if (not !found) && String.sub s i len_sub = sub then found := true\n done;\n !found\n\nlet ends_with s suffix = String.ends_with ~suffix s\n\nlet render ?(width = 400.) ?(height = 300.) spec =\n let tmp = Filename.temp_file \"hugin_test\" \".svg\" in\n Hugin.render_svg ~width ~height tmp spec;\n let ic = open_in tmp in\n let n = in_channel_length ic in\n let s = really_input_string ic n in\n close_in ic;\n Sys.remove tmp;\n s\n\nlet sample_x = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0))\nlet sample_y = Nx.init Float32 [| 5 |] (fun i -> float_of_int i.(0))\n\n(* SVG structure *)\n\nlet test_svg_envelope () =\n let svg = render (Hugin.line ~x:sample_x ~y:sample_y ()) in\n is_true ~msg:\"starts with xml\"\n (String.length svg > 5 && String.sub svg 0 5 = \"\\n\")\n\nlet test_svg_dimensions () =\n let svg =\n render ~width:800. ~height:600. (Hugin.line ~x:sample_x ~y:sample_y ())\n in\n is_true ~msg:\"has width\" (contains svg \"width=\\\"800\\\"\");\n is_true ~msg:\"has height\" (contains svg \"height=\\\"600\\\"\")\n\n(* XML escaping through text marks *)\n\nlet test_xml_escaping () =\n let svg = render (Hugin.text ~x:1. ~y:1. \"a & b < c\" ()) in\n is_true ~msg:\"ampersand escaped\" (contains svg \"&\");\n is_true ~msg:\"less-than escaped\" (contains svg \"<\")\n\nlet test_xml_escaping_quotes () =\n let svg = render (Hugin.text ~x:1. ~y:1. \"say \\\"hello\\\"\" ()) in\n is_true ~msg:\"quotes escaped\" (contains svg \""\")\n\n(* Clip regions *)\n\nlet test_clip_region () =\n (* A line plot should produce a clip region for the data area *)\n let svg = render (Hugin.line ~x:sample_x ~y:sample_y ()) in\n is_true ~msg:\"has clipPath\" (contains svg \"... pattern *)\n is_true ~msg:\"has use elements\" (contains svg \" len_s then false\n else\n let found = ref false in\n for i = 0 to len_s - len_sub do\n if (not !found) && String.sub s i len_sub = sub then found := true\n done;\n !found\n\nlet count_substring s sub =\n let len_s = String.length s and len_sub = String.length sub in\n if len_sub > len_s || len_sub = 0 then 0\n else begin\n let count = ref 0 in\n for i = 0 to len_s - len_sub do\n if String.sub s i len_sub = sub then incr count\n done;\n !count\n end\n\nlet render spec =\n let tmp = Filename.temp_file \"hugin_test\" \".svg\" in\n Hugin.render_svg ~width:400. ~height:300. tmp spec;\n let ic = open_in tmp in\n let n = in_channel_length ic in\n let s = really_input_string ic n in\n close_in ic;\n Sys.remove tmp;\n s\n\n(* linear tick formatting *)\n\nlet test_zero_label () =\n let x = Nx.create Float32 [| 5 |] [| -10.; -5.; 0.; 5.; 10. |] in\n let y = Nx.create Float32 [| 5 |] [| -10.; -5.; 0.; 5.; 10. |] in\n let svg = render (Hugin.line ~x ~y ()) in\n (* The zero tick should show \"0\" not \"1e-15\" or similar *)\n is_true ~msg:\"has zero tick\" (contains svg \">0<\")\n\nlet test_reasonable_count () =\n let x = Nx.init Float32 [| 101 |] (fun i -> float_of_int i.(0)) in\n let y = Nx.init Float32 [| 101 |] (fun i -> float_of_int i.(0)) in\n let svg = render (Hugin.line ~x ~y ()) in\n (* Count text elements that look like tick labels. Each tick generates a\n element. A basic line plot should have title-area text + x ticks + y\n ticks. We just check it's not an absurd number. *)\n let text_count = count_substring svg \" 2 && text_count < 40)\n\n(* log tick formatting *)\n\nlet test_log_tick_labels () =\n let x = Nx.create Float32 [| 5 |] [| 0.01; 0.1; 1.; 10.; 100. |] in\n let y = Nx.create Float32 [| 5 |] [| 1.; 2.; 3.; 4.; 5. |] in\n let svg = render (Hugin.line ~x ~y () |> Hugin.xscale `Log) in\n (* Log ticks should be powers of 10, formatted as 10^k *)\n is_true ~msg:\"has 10^-2\" (contains svg \">10^-2<\");\n is_true ~msg:\"has 10^2\" (contains svg \">10^2<\")\n\n(* large range doesn't explode *)\n\nlet test_large_range () =\n let x = Nx.create Float32 [| 2 |] [| 0.; 1e6 |] in\n let y = Nx.create Float32 [| 2 |] [| 0.; 1. |] in\n let svg = render (Hugin.line ~x ~y ()) in\n let text_count = count_substring svg \" Lexing.from_string\n |> !Toploop.parse_toplevel_phrase\n in\n Toploop.execute_phrase false Format.err_formatter phrase |> ignore\n\nlet () = install_printer \"Hugin.pp\"\n"
},
{
"path": "packages/hugin/ucairo/discover/discover.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule C = Configurator.V1\nmodule P = C.Pkg_config\n\nlet default_cairo c =\n let sys = C.ocaml_config_var_exn c \"system\" in\n if sys = \"msvc\" || sys = \"win64\" then\n {\n P.cflags = [ \"-I\"; \"C:\\\\gtk\\\\include\\\\cairo\" ];\n libs = [ \"/LC:\\\\gtk\\\\lib\"; \"cairo.lib\" ];\n }\n else { P.cflags = [ \"-I/usr/include/cairo\" ]; libs = [ \"-lcairo\" ] }\n\nlet () =\n C.main ~name:\"cairo-config\" (fun c ->\n let p =\n match P.get c with\n | Some p -> (\n match P.query p ~package:\"cairo\" with\n | Some p -> p\n | None -> default_cairo c)\n | None -> default_cairo c\n in\n let cflags =\n match Sys.getenv \"CAIRO_CFLAGS\" with\n | exception Not_found -> \"-fPIC\" :: p.P.cflags\n | alt -> C.Flags.extract_blank_separated_words alt\n in\n let libs =\n match Sys.getenv \"CAIRO_LIBS\" with\n | exception Not_found -> p.P.libs\n | alt -> C.Flags.extract_blank_separated_words alt\n in\n C.Flags.write_sexp \"cflags.sexp\" cflags;\n C.Flags.write_sexp \"clibs.sexp\" libs)\n"
},
{
"path": "packages/hugin/ucairo/discover/dune",
"content": "(executable\n (name discover)\n (modules discover)\n (libraries dune-configurator))\n\n(rule\n (deps discover.exe)\n (targets cflags.sexp clibs.sexp)\n (action\n (run ./discover.exe)))\n"
},
{
"path": "packages/hugin/ucairo/dune",
"content": "(library\n (name ucairo)\n (public_name hugin.ucairo)\n (foreign_stubs\n (language c)\n (names ucairo_stubs)\n (flags\n (:include discover/cflags.sexp)))\n (c_library_flags\n (:include discover/clibs.sexp)))\n"
},
{
"path": "packages/hugin/ucairo/ucairo.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Handle types *)\n\ntype t\ntype surface\n\n(* Enums *)\n\ntype font_weight = Normal | Bold\ntype line_cap = Butt | Round | Square\ntype line_join = Join_miter | Join_round | Join_bevel\n\ntype antialias =\n | Antialias_default\n | Antialias_none\n | Antialias_gray\n | Antialias_subpixel\n\n(* Text extents *)\n\ntype text_extents = {\n x_bearing : float;\n y_bearing : float;\n width : float;\n height : float;\n x_advance : float;\n y_advance : float;\n}\n\n(* Context *)\n\nexternal create : surface -> t = \"caml_ucairo_create\"\n\n(* State *)\n\nexternal save : t -> unit = \"caml_ucairo_save\"\nexternal restore : t -> unit = \"caml_ucairo_restore\"\n\n(* Transformations *)\n\nexternal translate : t -> float -> float -> unit = \"caml_ucairo_translate\"\nexternal scale : t -> float -> float -> unit = \"caml_ucairo_scale\"\nexternal rotate : t -> float -> unit = \"caml_ucairo_rotate\"\n\n(* Source *)\n\nexternal set_source_rgba : t -> float -> float -> float -> float -> unit\n = \"caml_ucairo_set_source_rgba\"\n\nexternal set_source_surface : t -> surface -> x:float -> y:float -> unit\n = \"caml_ucairo_set_source_surface\"\n\n(* Stroke/fill parameters *)\n\nexternal set_line_width : t -> float -> unit = \"caml_ucairo_set_line_width\"\nexternal raw_set_line_cap : t -> int -> unit = \"caml_ucairo_set_line_cap\"\n\nlet set_line_cap t cap =\n raw_set_line_cap t (match cap with Butt -> 0 | Round -> 1 | Square -> 2)\n\nexternal raw_set_line_join : t -> int -> unit = \"caml_ucairo_set_line_join\"\n\nlet set_line_join t join =\n raw_set_line_join t\n (match join with Join_miter -> 0 | Join_round -> 1 | Join_bevel -> 2)\n\nexternal set_dash : t -> float array -> unit = \"caml_ucairo_set_dash\"\nexternal raw_set_antialias : t -> int -> unit = \"caml_ucairo_set_antialias\"\n\nlet set_antialias t aa =\n raw_set_antialias t\n (match aa with\n | Antialias_default -> 0\n | Antialias_none -> 1\n | Antialias_gray -> 2\n | Antialias_subpixel -> 3)\n\n(* Font *)\n\nexternal raw_select_font_face : t -> string -> int -> unit\n = \"caml_ucairo_select_font_face\"\n\nlet select_font_face t family weight =\n raw_select_font_face t family (match weight with Normal -> 0 | Bold -> 1)\n\nexternal set_font_size : t -> float -> unit = \"caml_ucairo_set_font_size\"\nexternal text_extents : t -> string -> text_extents = \"caml_ucairo_text_extents\"\nexternal show_text : t -> string -> unit = \"caml_ucairo_show_text\"\n\n(* Path *)\n\nexternal move_to : t -> float -> float -> unit = \"caml_ucairo_move_to\"\nexternal line_to : t -> float -> float -> unit = \"caml_ucairo_line_to\"\n\nexternal arc : t -> float -> float -> r:float -> a1:float -> a2:float -> unit\n = \"caml_ucairo_arc_bytecode\" \"caml_ucairo_arc_native\"\n\nexternal rectangle : t -> float -> float -> w:float -> h:float -> unit\n = \"caml_ucairo_rectangle\"\n\nmodule Path = struct\n external close : t -> unit = \"caml_ucairo_path_close\"\n external clear : t -> unit = \"caml_ucairo_path_clear\"\nend\n\n(* Drawing *)\n\nexternal fill : t -> unit = \"caml_ucairo_fill\"\nexternal fill_preserve : t -> unit = \"caml_ucairo_fill_preserve\"\nexternal stroke : t -> unit = \"caml_ucairo_stroke\"\nexternal paint : t -> unit = \"caml_ucairo_paint\"\nexternal clip : t -> unit = \"caml_ucairo_clip\"\n\n(* Surface *)\n\nmodule Surface = struct\n external finish : surface -> unit = \"caml_ucairo_surface_finish\"\n external flush : surface -> unit = \"caml_ucairo_surface_flush\"\nend\n\n(* Image *)\n\nmodule Image = struct\n external create : w:int -> h:int -> surface = \"caml_ucairo_image_create\"\n\n external create_for_data8 :\n (int, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t ->\n w:int ->\n h:int ->\n stride:int ->\n surface = \"caml_ucairo_image_create_for_data8\"\n\n external stride_for_width : int -> int = \"caml_ucairo_image_stride_for_width\"\n [@@noalloc]\nend\n\n(* PDF *)\n\nmodule Pdf = struct\n external create : string -> w:float -> h:float -> surface\n = \"caml_ucairo_pdf_create\"\nend\n\n(* PNG *)\n\nmodule Png = struct\n external write : surface -> string -> unit = \"caml_ucairo_png_write\"\n\n external write_to_stream : surface -> (string -> unit) -> unit\n = \"caml_ucairo_png_write_to_stream\"\nend\n"
},
{
"path": "packages/hugin/ucairo/ucairo.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Minimal Cairo bindings.\n\n Thin bindings covering image and PDF surface creation, path drawing, text\n rendering, and PNG output. Designed for the Hugin rendering backend; not a\n general-purpose Cairo binding.\n\n All functions raise [Failure] on Cairo errors and [Invalid_argument] on\n destroyed handles. *)\n\n(** {1:types Handle types} *)\n\ntype t\n(** The type for Cairo drawing contexts. *)\n\ntype surface\n(** The type for Cairo surfaces. *)\n\n(** {1:enums Enumerations} *)\n\ntype font_weight = Normal | Bold (** The type for font weight. *)\ntype line_cap = Butt | Round | Square (** The type for line cap style. *)\n\ntype line_join =\n | Join_miter\n | Join_round\n | Join_bevel (** The type for line join style. *)\n\ntype antialias =\n | Antialias_default\n | Antialias_none\n | Antialias_gray\n | Antialias_subpixel (** The type for antialiasing mode. *)\n\n(** {1:text_extents Text extents} *)\n\ntype text_extents = {\n x_bearing : float;\n y_bearing : float;\n width : float;\n height : float;\n x_advance : float;\n y_advance : float;\n}\n(** The type for text extent measurements. *)\n\n(** {1:context Context creation} *)\n\nval create : surface -> t\n(** [create surface] is a new drawing context targeting [surface]. *)\n\n(** {1:state State} *)\n\nval save : t -> unit\n(** [save t] pushes the current graphics state onto the stack. *)\n\nval restore : t -> unit\n(** [restore t] pops the graphics state from the stack. *)\n\n(** {1:transform Transformations} *)\n\nval translate : t -> float -> float -> unit\n(** [translate t tx ty] translates the user-space origin by [(tx, ty)]. *)\n\nval scale : t -> float -> float -> unit\n(** [scale t sx sy] scales the user-space axes by [(sx, sy)]. *)\n\nval rotate : t -> float -> unit\n(** [rotate t angle] rotates the user-space axes by [angle] radians. *)\n\n(** {1:source Source} *)\n\nval set_source_rgba : t -> float -> float -> float -> float -> unit\n(** [set_source_rgba t r g b a] sets the source to the given RGBA color. *)\n\nval set_source_surface : t -> surface -> x:float -> y:float -> unit\n(** [set_source_surface t s ~x ~y] sets [s] as the source, offset by [(x, y)].\n*)\n\n(** {1:stroke_fill Stroke and fill parameters} *)\n\nval set_line_width : t -> float -> unit\n(** [set_line_width t w] sets the stroke line width. *)\n\nval set_line_cap : t -> line_cap -> unit\n(** [set_line_cap t cap] sets the line cap style. *)\n\nval set_line_join : t -> line_join -> unit\n(** [set_line_join t join] sets the line join style. *)\n\nval set_dash : t -> float array -> unit\n(** [set_dash t dashes] sets the dash pattern. An empty array disables dashing.\n*)\n\nval set_antialias : t -> antialias -> unit\n(** [set_antialias t aa] sets the antialiasing mode. *)\n\n(** {1:font Font} *)\n\nval select_font_face : t -> string -> font_weight -> unit\n(** [select_font_face t family weight] selects a toy font face. Slant is always\n upright. *)\n\nval set_font_size : t -> float -> unit\n(** [set_font_size t size] sets the font size in user-space units. *)\n\nval text_extents : t -> string -> text_extents\n(** [text_extents t s] is the extents of [s] with the current font. *)\n\nval show_text : t -> string -> unit\n(** [show_text t s] renders [s] at the current point. *)\n\n(** {1:path Path operations} *)\n\nval move_to : t -> float -> float -> unit\n(** [move_to t x y] begins a new sub-path at [(x, y)]. *)\n\nval line_to : t -> float -> float -> unit\n(** [line_to t x y] adds a line segment to [(x, y)]. *)\n\nval arc : t -> float -> float -> r:float -> a1:float -> a2:float -> unit\n(** [arc t xc yc ~r ~a1 ~a2] adds a circular arc centered at [(xc, yc)] with\n radius [r] from angle [a1] to [a2] (in radians). *)\n\nval rectangle : t -> float -> float -> w:float -> h:float -> unit\n(** [rectangle t x y ~w ~h] adds a closed rectangle sub-path. *)\n\n(** {1:path_mod Path module} *)\n\nmodule Path : sig\n val close : t -> unit\n (** [close t] closes the current sub-path with a line to its start. *)\n\n val clear : t -> unit\n (** [clear t] clears the current path. *)\nend\n\n(** {1:drawing Drawing operations} *)\n\nval fill : t -> unit\n(** [fill t] fills the current path and clears it. *)\n\nval fill_preserve : t -> unit\n(** [fill_preserve t] fills the current path without clearing it. *)\n\nval stroke : t -> unit\n(** [stroke t] strokes the current path and clears it. *)\n\nval paint : t -> unit\n(** [paint t] paints the current source everywhere within the current clip. *)\n\nval clip : t -> unit\n(** [clip t] establishes a new clip region by intersecting the current clip with\n the current path, then clears the path. *)\n\n(** {1:surface Surface operations} *)\n\nmodule Surface : sig\n val finish : surface -> unit\n (** [finish s] finalizes the surface and releases external resources. *)\n\n val flush : surface -> unit\n (** [flush s] completes any pending drawing operations. *)\nend\n\n(** {1:image Image surface} *)\n\nmodule Image : sig\n val create : w:int -> h:int -> surface\n (** [create ~w ~h] is a new ARGB32 image surface of dimensions [w] x [h].\n\n Raises [Failure] if allocation fails. *)\n\n val create_for_data8 :\n (int, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t ->\n w:int ->\n h:int ->\n stride:int ->\n surface\n (** [create_for_data8 data ~w ~h ~stride] wraps existing pixel [data] as an\n ARGB32 image surface. [data] must remain live for the lifetime of the\n surface. *)\n\n val stride_for_width : int -> int\n (** [stride_for_width w] is the minimum stride in bytes for an ARGB32 image of\n width [w], respecting Cairo alignment requirements. *)\nend\n\n(** {1:pdf PDF surface} *)\n\nmodule Pdf : sig\n val create : string -> w:float -> h:float -> surface\n (** [create filename ~w ~h] is a new PDF surface writing to [filename].\n Dimensions are in points (1 point = 1/72 inch). *)\nend\n\n(** {1:png PNG output} *)\n\nmodule Png : sig\n val write : surface -> string -> unit\n (** [write surface filename] writes [surface] as a PNG file. *)\n\n val write_to_stream : surface -> (string -> unit) -> unit\n (** [write_to_stream surface f] writes [surface] as PNG data, calling [f] with\n each chunk. *)\nend\n"
},
{
"path": "packages/hugin/ucairo/ucairo_stubs.c",
"content": "/*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*/\n\n#define CAML_NAME_SPACE\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \n#include \n#include \n\n/* --------------------------------------------------------------------------\n Custom blocks\n -------------------------------------------------------------------------- */\n\n/* Context (cairo_t *) */\n\n#define Context_val(v) (*(cairo_t **)Data_custom_val(v))\n\nstatic void finalize_context(value v)\n{\n cairo_t *cr = Context_val(v);\n if (cr != NULL) { cairo_destroy(cr); Context_val(v) = NULL; }\n}\n\nstatic struct custom_operations context_ops = {\n .identifier = \"ucairo.context\",\n .finalize = finalize_context,\n .compare = custom_compare_default,\n .hash = custom_hash_default,\n .serialize = custom_serialize_default,\n .deserialize = custom_deserialize_default,\n .compare_ext = custom_compare_ext_default,\n};\n\nstatic value alloc_context(cairo_t *cr)\n{\n value v = caml_alloc_custom(&context_ops, sizeof(cairo_t *), 0, 1);\n Context_val(v) = cr;\n return v;\n}\n\n/* Surface (cairo_surface_t *) */\n\n#define Surface_val(v) (*(cairo_surface_t **)Data_custom_val(v))\n\nstatic void finalize_surface(value v)\n{\n cairo_surface_t *s = Surface_val(v);\n if (s != NULL) { cairo_surface_destroy(s); Surface_val(v) = NULL; }\n}\n\nstatic struct custom_operations surface_ops = {\n .identifier = \"ucairo.surface\",\n .finalize = finalize_surface,\n .compare = custom_compare_default,\n .hash = custom_hash_default,\n .serialize = custom_serialize_default,\n .deserialize = custom_deserialize_default,\n .compare_ext = custom_compare_ext_default,\n};\n\nstatic value alloc_surface(cairo_surface_t *s)\n{\n value v = caml_alloc_custom(&surface_ops, sizeof(cairo_surface_t *), 0, 1);\n Surface_val(v) = s;\n return v;\n}\n\n/* --------------------------------------------------------------------------\n Helpers\n -------------------------------------------------------------------------- */\n\nstatic inline cairo_t *check_context(value v, const char *fn)\n{\n cairo_t *cr = Context_val(v);\n if (cr == NULL) caml_invalid_argument(fn);\n return cr;\n}\n\nstatic inline cairo_surface_t *check_surface(value v, const char *fn)\n{\n cairo_surface_t *s = Surface_val(v);\n if (s == NULL) caml_invalid_argument(fn);\n return s;\n}\n\n/* --------------------------------------------------------------------------\n Context creation\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_create(value vsurf)\n{\n CAMLparam1(vsurf);\n cairo_surface_t *s = check_surface(vsurf, \"Ucairo.create: destroyed surface\");\n cairo_t *cr = cairo_create(s);\n if (cairo_status(cr) != CAIRO_STATUS_SUCCESS) {\n const char *msg = cairo_status_to_string(cairo_status(cr));\n cairo_destroy(cr);\n caml_failwith(msg);\n }\n CAMLreturn(alloc_context(cr));\n}\n\n/* --------------------------------------------------------------------------\n State\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_save(value vcr)\n{\n CAMLparam1(vcr);\n cairo_save(check_context(vcr, \"Ucairo.save: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_restore(value vcr)\n{\n CAMLparam1(vcr);\n cairo_restore(check_context(vcr, \"Ucairo.restore: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\n/* --------------------------------------------------------------------------\n Transformations\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_translate(value vcr, value vtx, value vty)\n{\n CAMLparam3(vcr, vtx, vty);\n cairo_translate(check_context(vcr, \"Ucairo.translate: destroyed context\"),\n Double_val(vtx), Double_val(vty));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_scale(value vcr, value vsx, value vsy)\n{\n CAMLparam3(vcr, vsx, vsy);\n cairo_scale(check_context(vcr, \"Ucairo.scale: destroyed context\"),\n Double_val(vsx), Double_val(vsy));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_rotate(value vcr, value vangle)\n{\n CAMLparam2(vcr, vangle);\n cairo_rotate(check_context(vcr, \"Ucairo.rotate: destroyed context\"),\n Double_val(vangle));\n CAMLreturn(Val_unit);\n}\n\n/* --------------------------------------------------------------------------\n Source\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_set_source_rgba(value vcr, value vr, value vg,\n value vb, value va)\n{\n CAMLparam5(vcr, vr, vg, vb, va);\n cairo_set_source_rgba(check_context(vcr, \"Ucairo.set_source_rgba: destroyed context\"),\n Double_val(vr), Double_val(vg),\n Double_val(vb), Double_val(va));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_set_source_surface(value vcr, value vsurf,\n value vx, value vy)\n{\n CAMLparam4(vcr, vsurf, vx, vy);\n cairo_set_source_surface(\n check_context(vcr, \"Ucairo.set_source_surface: destroyed context\"),\n check_surface(vsurf, \"Ucairo.set_source_surface: destroyed surface\"),\n Double_val(vx), Double_val(vy));\n CAMLreturn(Val_unit);\n}\n\n/* --------------------------------------------------------------------------\n Stroke and fill parameters\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_set_line_width(value vcr, value vw)\n{\n CAMLparam2(vcr, vw);\n cairo_set_line_width(check_context(vcr, \"Ucairo.set_line_width: destroyed context\"),\n Double_val(vw));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_set_line_cap(value vcr, value vcap)\n{\n CAMLparam2(vcr, vcap);\n static const cairo_line_cap_t caps[] = {\n CAIRO_LINE_CAP_BUTT, CAIRO_LINE_CAP_ROUND, CAIRO_LINE_CAP_SQUARE\n };\n cairo_set_line_cap(check_context(vcr, \"Ucairo.set_line_cap: destroyed context\"),\n caps[Int_val(vcap)]);\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_set_line_join(value vcr, value vjoin)\n{\n CAMLparam2(vcr, vjoin);\n static const cairo_line_join_t joins[] = {\n CAIRO_LINE_JOIN_MITER, CAIRO_LINE_JOIN_ROUND, CAIRO_LINE_JOIN_BEVEL\n };\n cairo_set_line_join(check_context(vcr, \"Ucairo.set_line_join: destroyed context\"),\n joins[Int_val(vjoin)]);\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_set_dash(value vcr, value varr)\n{\n CAMLparam2(vcr, varr);\n cairo_t *cr = check_context(vcr, \"Ucairo.set_dash: destroyed context\");\n int n = Wosize_val(varr) / Double_wosize;\n if (n == 0) {\n cairo_set_dash(cr, NULL, 0, 0.0);\n } else {\n double stack_buf[64];\n double *dashes = n <= 64 ? stack_buf\n : caml_stat_alloc(n * sizeof(double));\n for (int i = 0; i < n; i++)\n dashes[i] = Double_field(varr, i);\n cairo_set_dash(cr, dashes, n, 0.0);\n if (dashes != stack_buf) caml_stat_free(dashes);\n }\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_set_antialias(value vcr, value vaa)\n{\n CAMLparam2(vcr, vaa);\n static const cairo_antialias_t aa[] = {\n CAIRO_ANTIALIAS_DEFAULT, CAIRO_ANTIALIAS_NONE,\n CAIRO_ANTIALIAS_GRAY, CAIRO_ANTIALIAS_SUBPIXEL\n };\n cairo_set_antialias(check_context(vcr, \"Ucairo.set_antialias: destroyed context\"),\n aa[Int_val(vaa)]);\n CAMLreturn(Val_unit);\n}\n\n/* --------------------------------------------------------------------------\n Font\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_select_font_face(value vcr, value vfamily,\n value vweight)\n{\n CAMLparam3(vcr, vfamily, vweight);\n static const cairo_font_weight_t weights[] = {\n CAIRO_FONT_WEIGHT_NORMAL, CAIRO_FONT_WEIGHT_BOLD\n };\n cairo_select_font_face(\n check_context(vcr, \"Ucairo.select_font_face: destroyed context\"),\n String_val(vfamily), CAIRO_FONT_SLANT_NORMAL, weights[Int_val(vweight)]);\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_set_font_size(value vcr, value vsize)\n{\n CAMLparam2(vcr, vsize);\n cairo_set_font_size(check_context(vcr, \"Ucairo.set_font_size: destroyed context\"),\n Double_val(vsize));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_text_extents(value vcr, value vstr)\n{\n CAMLparam2(vcr, vstr);\n CAMLlocal1(result);\n cairo_t *cr = check_context(vcr, \"Ucairo.text_extents: destroyed context\");\n cairo_text_extents_t ext;\n cairo_text_extents(cr, String_val(vstr), &ext);\n result = caml_alloc(6 * Double_wosize, Double_array_tag);\n Store_double_field(result, 0, ext.x_bearing);\n Store_double_field(result, 1, ext.y_bearing);\n Store_double_field(result, 2, ext.width);\n Store_double_field(result, 3, ext.height);\n Store_double_field(result, 4, ext.x_advance);\n Store_double_field(result, 5, ext.y_advance);\n CAMLreturn(result);\n}\n\nCAMLprim value caml_ucairo_show_text(value vcr, value vstr)\n{\n CAMLparam2(vcr, vstr);\n cairo_show_text(check_context(vcr, \"Ucairo.show_text: destroyed context\"),\n String_val(vstr));\n CAMLreturn(Val_unit);\n}\n\n/* --------------------------------------------------------------------------\n Path operations\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_move_to(value vcr, value vx, value vy)\n{\n CAMLparam3(vcr, vx, vy);\n cairo_move_to(check_context(vcr, \"Ucairo.move_to: destroyed context\"),\n Double_val(vx), Double_val(vy));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_line_to(value vcr, value vx, value vy)\n{\n CAMLparam3(vcr, vx, vy);\n cairo_line_to(check_context(vcr, \"Ucairo.line_to: destroyed context\"),\n Double_val(vx), Double_val(vy));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_arc_native(value vcr, value vxc, value vyc,\n value vr, value va1, value va2)\n{\n CAMLparam5(vcr, vxc, vyc, vr, va1);\n CAMLxparam1(va2);\n cairo_arc(check_context(vcr, \"Ucairo.arc: destroyed context\"),\n Double_val(vxc), Double_val(vyc),\n Double_val(vr), Double_val(va1), Double_val(va2));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_arc_bytecode(value *argv, int argc)\n{\n (void)argc;\n return caml_ucairo_arc_native(argv[0], argv[1], argv[2],\n argv[3], argv[4], argv[5]);\n}\n\nCAMLprim value caml_ucairo_rectangle(value vcr, value vx, value vy,\n value vw, value vh)\n{\n CAMLparam5(vcr, vx, vy, vw, vh);\n cairo_rectangle(check_context(vcr, \"Ucairo.rectangle: destroyed context\"),\n Double_val(vx), Double_val(vy),\n Double_val(vw), Double_val(vh));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_path_close(value vcr)\n{\n CAMLparam1(vcr);\n cairo_close_path(check_context(vcr, \"Ucairo.Path.close: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_path_clear(value vcr)\n{\n CAMLparam1(vcr);\n cairo_new_path(check_context(vcr, \"Ucairo.Path.clear: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\n/* --------------------------------------------------------------------------\n Drawing operations\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_fill(value vcr)\n{\n CAMLparam1(vcr);\n cairo_fill(check_context(vcr, \"Ucairo.fill: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_fill_preserve(value vcr)\n{\n CAMLparam1(vcr);\n cairo_fill_preserve(check_context(vcr, \"Ucairo.fill_preserve: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_stroke(value vcr)\n{\n CAMLparam1(vcr);\n cairo_stroke(check_context(vcr, \"Ucairo.stroke: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_paint(value vcr)\n{\n CAMLparam1(vcr);\n cairo_paint(check_context(vcr, \"Ucairo.paint: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_clip(value vcr)\n{\n CAMLparam1(vcr);\n cairo_clip(check_context(vcr, \"Ucairo.clip: destroyed context\"));\n CAMLreturn(Val_unit);\n}\n\n/* --------------------------------------------------------------------------\n Surface operations\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_surface_finish(value vsurf)\n{\n CAMLparam1(vsurf);\n cairo_surface_t *s = Surface_val(vsurf);\n if (s != NULL) cairo_surface_finish(s);\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_ucairo_surface_flush(value vsurf)\n{\n CAMLparam1(vsurf);\n cairo_surface_flush(check_surface(vsurf, \"Ucairo.Surface.flush: destroyed surface\"));\n CAMLreturn(Val_unit);\n}\n\n/* --------------------------------------------------------------------------\n Image surface\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_image_create(value vw, value vh)\n{\n CAMLparam2(vw, vh);\n cairo_surface_t *s = cairo_image_surface_create(\n CAIRO_FORMAT_ARGB32, Int_val(vw), Int_val(vh));\n if (cairo_surface_status(s) != CAIRO_STATUS_SUCCESS) {\n const char *msg = cairo_status_to_string(cairo_surface_status(s));\n cairo_surface_destroy(s);\n caml_failwith(msg);\n }\n CAMLreturn(alloc_surface(s));\n}\n\nCAMLprim value caml_ucairo_image_create_for_data8(value vdata, value vw,\n value vh, value vstride)\n{\n CAMLparam4(vdata, vw, vh, vstride);\n unsigned char *data = (unsigned char *)Caml_ba_data_val(vdata);\n cairo_surface_t *s = cairo_image_surface_create_for_data(\n data, CAIRO_FORMAT_ARGB32, Int_val(vw), Int_val(vh), Int_val(vstride));\n if (cairo_surface_status(s) != CAIRO_STATUS_SUCCESS) {\n const char *msg = cairo_status_to_string(cairo_surface_status(s));\n cairo_surface_destroy(s);\n caml_failwith(msg);\n }\n CAMLreturn(alloc_surface(s));\n}\n\nCAMLprim value caml_ucairo_image_stride_for_width(value vw)\n{\n return Val_int(cairo_format_stride_for_width(CAIRO_FORMAT_ARGB32, Int_val(vw)));\n}\n\n/* --------------------------------------------------------------------------\n PDF surface\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_pdf_create(value vfilename, value vw, value vh)\n{\n CAMLparam3(vfilename, vw, vh);\n cairo_surface_t *s = cairo_pdf_surface_create(\n String_val(vfilename), Double_val(vw), Double_val(vh));\n if (cairo_surface_status(s) != CAIRO_STATUS_SUCCESS) {\n const char *msg = cairo_status_to_string(cairo_surface_status(s));\n cairo_surface_destroy(s);\n caml_failwith(msg);\n }\n CAMLreturn(alloc_surface(s));\n}\n\n/* --------------------------------------------------------------------------\n PNG output\n -------------------------------------------------------------------------- */\n\nCAMLprim value caml_ucairo_png_write(value vsurf, value vfilename)\n{\n CAMLparam2(vsurf, vfilename);\n cairo_surface_t *s = check_surface(vsurf, \"Ucairo.Png.write: destroyed surface\");\n cairo_status_t st = cairo_surface_write_to_png(s, String_val(vfilename));\n if (st != CAIRO_STATUS_SUCCESS)\n caml_failwith(cairo_status_to_string(st));\n CAMLreturn(Val_unit);\n}\n\nstatic cairo_status_t\npng_write_func(void *closure, const unsigned char *data, unsigned int length)\n{\n CAMLparam0();\n CAMLlocal2(vstr, r);\n vstr = caml_alloc_string(length);\n memcpy(Bytes_val(vstr), data, length);\n /* closure points to CAMLparam-rooted vcallback in the caller frame;\n re-read after allocation so we get the post-GC value. */\n r = caml_callback_exn(*(value *)closure, vstr);\n if (Is_exception_result(r))\n CAMLreturnT(cairo_status_t, CAIRO_STATUS_WRITE_ERROR);\n CAMLreturnT(cairo_status_t, CAIRO_STATUS_SUCCESS);\n}\n\nCAMLprim value caml_ucairo_png_write_to_stream(value vsurf, value vcallback)\n{\n CAMLparam2(vsurf, vcallback);\n cairo_surface_t *s = check_surface(vsurf,\n \"Ucairo.Png.write_to_stream: destroyed surface\");\n /* vcallback is rooted by CAMLparam2; we pass its address as the closure\n so the callback can retrieve it. This is safe because\n cairo_surface_write_to_png_stream calls png_write_func synchronously. */\n cairo_status_t st = cairo_surface_write_to_png_stream(\n s, png_write_func, &vcallback);\n if (st != CAIRO_STATUS_SUCCESS)\n caml_failwith(cairo_status_to_string(st));\n CAMLreturn(Val_unit);\n}\n"
},
{
"path": "packages/hugin/usdl/discover/discover.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule C = Configurator.V1\n\nlet () =\n C.main ~name:\"sdl2-config\" (fun c ->\n let pkg_config =\n match C.Pkg_config.get c with\n | None -> C.die \"pkg-config not found\"\n | Some pc -> pc\n in\n let sdl2 =\n match C.Pkg_config.query pkg_config ~package:\"sdl2\" with\n | None -> C.die \"SDL2 not found via pkg-config\"\n | Some info -> info\n in\n let cflags = \"-fPIC\" :: sdl2.C.Pkg_config.cflags in\n let clibs = sdl2.C.Pkg_config.libs in\n C.Flags.write_sexp \"cflags.sexp\" cflags;\n C.Flags.write_sexp \"clibs.sexp\" clibs)\n"
},
{
"path": "packages/hugin/usdl/discover/dune",
"content": "(executable\n (name discover)\n (modules discover)\n (libraries dune-configurator))\n\n(rule\n (deps discover.exe)\n (targets cflags.sexp clibs.sexp)\n (action\n (run ./discover.exe)))\n"
},
{
"path": "packages/hugin/usdl/dune",
"content": "(library\n (name usdl)\n (public_name hugin.usdl)\n (foreign_stubs\n (language c)\n (names usdl_stubs)\n (flags\n (:include discover/cflags.sexp)))\n (c_library_flags\n (:include discover/clibs.sexp)))\n"
},
{
"path": "packages/hugin/usdl/usdl.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Init / quit *)\n\nexternal init : unit -> unit = \"caml_usdl_init\"\nexternal quit : unit -> unit = \"caml_usdl_quit\"\n\n(* Handle types *)\n\ntype renderer\ntype surface\ntype texture\n\n(* Window *)\n\nmodule Window = struct\n type t\n\n external create : title:string -> w:int -> h:int -> t\n = \"caml_usdl_window_create\"\n\n external destroy : t -> unit = \"caml_usdl_window_destroy\"\nend\n\n(* Renderer *)\n\nmodule Renderer = struct\n type t = renderer\n\n external create : Window.t -> t = \"caml_usdl_renderer_create\"\n external output_size : t -> int * int = \"caml_usdl_renderer_output_size\"\n external clear : t -> unit = \"caml_usdl_renderer_clear\"\n external copy : t -> texture -> unit = \"caml_usdl_renderer_copy\"\n external present : t -> unit = \"caml_usdl_renderer_present\"\n external destroy : t -> unit = \"caml_usdl_renderer_destroy\"\nend\n\n(* Surface *)\n\nmodule Surface = struct\n type t = surface\n\n external create_argb8888 : w:int -> h:int -> t\n = \"caml_usdl_surface_create_argb8888\"\n\n external pitch : t -> int = \"caml_usdl_surface_pitch\"\n\n external pixels :\n t -> (int, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t\n = \"caml_usdl_surface_pixels\"\n\n external destroy : t -> unit = \"caml_usdl_surface_destroy\"\nend\n\n(* Texture *)\n\nmodule Texture = struct\n type t = texture\n\n external of_surface : Renderer.t -> Surface.t -> t\n = \"caml_usdl_texture_of_surface\"\n\n external destroy : t -> unit = \"caml_usdl_texture_destroy\"\nend\n\n(* Event *)\n\nmodule Event = struct\n type t\n type event_type = [ `Quit | `Window_event | `Key_down | `Unknown of int ]\n\n type window_event =\n [ `Resized | `Size_changed | `Exposed | `Close | `Unknown of int ]\n\n external create : unit -> t = \"caml_usdl_event_create\"\n external wait : t -> bool = \"caml_usdl_event_wait\"\n external raw_type : t -> int = \"caml_usdl_event_type\" [@@noalloc]\n external raw_window_id : t -> int = \"caml_usdl_event_window_id\" [@@noalloc]\n external keycode : t -> int = \"caml_usdl_event_keycode\" [@@noalloc]\n\n let typ t =\n match raw_type t with\n | 0x100 -> `Quit\n | 0x200 -> `Window_event\n | 0x300 -> `Key_down\n | n -> `Unknown n\n\n let window_event_id t =\n match raw_window_id t with\n | 5 -> `Resized\n | 6 -> `Size_changed\n | 2 -> `Exposed\n | 14 -> `Close\n | n -> `Unknown n\nend\n\n(* Keycodes *)\n\nmodule Keycode = struct\n let escape = 27\n let q = Char.code 'q'\nend\n"
},
{
"path": "packages/hugin/usdl/usdl.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Minimal SDL2 bindings.\n\n Thin bindings covering window creation, renderer management, surface pixel\n access, and event polling. Designed for the Cairo-SDL integration layer; not\n a general-purpose SDL binding.\n\n All functions raise [Failure] on SDL errors. *)\n\n(** {1:init Initialization} *)\n\nval init : unit -> unit\n(** [init ()] initializes SDL video and sets the render scale quality hint.\n\n Raises [Failure] if SDL initialization fails. *)\n\nval quit : unit -> unit\n(** [quit ()] shuts down SDL. *)\n\n(** {1:handles Handle types} *)\n\ntype renderer\n(** The type for SDL renderers. *)\n\ntype surface\n(** The type for SDL surfaces. *)\n\ntype texture\n(** The type for SDL textures. *)\n\n(** {1:window Window} *)\n\nmodule Window : sig\n type t\n (** The type for SDL windows. *)\n\n val create : title:string -> w:int -> h:int -> t\n (** [create ~title ~w ~h] is a resizable, high-DPI-aware window.\n\n Raises [Failure] if window creation fails. *)\n\n val destroy : t -> unit\n (** [destroy t] frees the window. Safe to call more than once. *)\nend\n\n(** {1:renderer Renderer} *)\n\nmodule Renderer : sig\n type t = renderer\n (** The type for SDL renderers. *)\n\n val create : Window.t -> t\n (** [create win] is a hardware-accelerated, vsync-enabled renderer for [win].\n\n Raises [Failure] if renderer creation fails. *)\n\n val output_size : t -> int * int\n (** [output_size t] is [(w, h)] in pixels (accounting for high-DPI scaling).\n\n Raises [Failure] if the query fails. *)\n\n val clear : t -> unit\n (** [clear t] clears the render target. *)\n\n val copy : t -> texture -> unit\n (** [copy t tex] copies [tex] to the entire render target. *)\n\n val present : t -> unit\n (** [present t] presents the composed backbuffer. *)\n\n val destroy : t -> unit\n (** [destroy t] frees the renderer. Safe to call more than once. *)\nend\n\n(** {1:surface Surface} *)\n\nmodule Surface : sig\n type t = surface\n (** The type for SDL surfaces. *)\n\n val create_argb8888 : w:int -> h:int -> t\n (** [create_argb8888 ~w ~h] is a 32-bit ARGB8888 surface.\n\n Raises [Failure] if allocation fails. *)\n\n val pitch : t -> int\n (** [pitch t] is the byte length of one row. *)\n\n val pixels :\n t -> (int, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t\n (** [pixels t] is the raw pixel buffer. The bigarray is a view onto\n SDL-managed memory; it must not outlive the surface. *)\n\n val destroy : t -> unit\n (** [destroy t] frees the surface. Safe to call more than once. *)\nend\n\n(** {1:texture Texture} *)\n\nmodule Texture : sig\n type t = texture\n (** The type for SDL textures. *)\n\n val of_surface : renderer -> surface -> t\n (** [of_surface ren surf] creates a texture from [surf].\n\n Raises [Failure] if texture creation fails. *)\n\n val destroy : t -> unit\n (** [destroy t] frees the texture. Safe to call more than once. *)\nend\n\n(** {1:event Events} *)\n\nmodule Event : sig\n type t\n (** The type for event storage. *)\n\n type event_type = [ `Quit | `Window_event | `Key_down | `Unknown of int ]\n (** The type for event kinds. *)\n\n type window_event =\n [ `Resized | `Size_changed | `Exposed | `Close | `Unknown of int ]\n (** The type for window event sub-kinds. *)\n\n val create : unit -> t\n (** [create ()] allocates event storage. *)\n\n val wait : t -> bool\n (** [wait t] blocks until an event arrives. Returns [true] if an event was\n received, [false] on error. Releases the runtime lock while blocking. *)\n\n val typ : t -> event_type\n (** [typ t] is the kind of the last received event. *)\n\n val window_event_id : t -> window_event\n (** [window_event_id t] is the window event sub-kind. Only meaningful when\n [typ t] is [`Window_event]. *)\n\n val keycode : t -> int\n (** [keycode t] is the key code. Only meaningful when [typ t] is [`Key_down].\n *)\nend\n\n(** {1:keycode Key codes} *)\n\nmodule Keycode : sig\n val escape : int\n val q : int\nend\n"
},
{
"path": "packages/hugin/usdl/usdl_stubs.c",
"content": "/*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*/\n\n#define CAML_NAME_SPACE\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#ifdef __APPLE__\n#include \n#else\n#include \n#endif\n\n/* Window */\n\n#define Window_val(v) (*(SDL_Window **)Data_custom_val(v))\n\nstatic void finalize_window(value v)\n{\n SDL_Window *w = Window_val(v);\n if (w != NULL) { SDL_DestroyWindow(w); Window_val(v) = NULL; }\n}\n\nstatic struct custom_operations window_ops = {\n .identifier = \"usdl.window\",\n .finalize = finalize_window,\n .compare = custom_compare_default,\n .hash = custom_hash_default,\n .serialize = custom_serialize_default,\n .deserialize = custom_deserialize_default,\n .compare_ext = custom_compare_ext_default,\n};\n\nstatic value alloc_window(SDL_Window *w)\n{\n value v = caml_alloc_custom(&window_ops, sizeof(SDL_Window *), 0, 1);\n Window_val(v) = w;\n return v;\n}\n\n/* Renderer */\n\n#define Renderer_val(v) (*(SDL_Renderer **)Data_custom_val(v))\n\nstatic void finalize_renderer(value v)\n{\n SDL_Renderer *r = Renderer_val(v);\n if (r != NULL) { SDL_DestroyRenderer(r); Renderer_val(v) = NULL; }\n}\n\nstatic struct custom_operations renderer_ops = {\n .identifier = \"usdl.renderer\",\n .finalize = finalize_renderer,\n .compare = custom_compare_default,\n .hash = custom_hash_default,\n .serialize = custom_serialize_default,\n .deserialize = custom_deserialize_default,\n .compare_ext = custom_compare_ext_default,\n};\n\nstatic value alloc_renderer(SDL_Renderer *r)\n{\n value v = caml_alloc_custom(&renderer_ops, sizeof(SDL_Renderer *), 0, 1);\n Renderer_val(v) = r;\n return v;\n}\n\n/* Surface */\n\n#define Surface_val(v) (*(SDL_Surface **)Data_custom_val(v))\n\nstatic void finalize_surface(value v)\n{\n SDL_Surface *s = Surface_val(v);\n if (s != NULL) { SDL_FreeSurface(s); Surface_val(v) = NULL; }\n}\n\nstatic struct custom_operations surface_ops = {\n .identifier = \"usdl.surface\",\n .finalize = finalize_surface,\n .compare = custom_compare_default,\n .hash = custom_hash_default,\n .serialize = custom_serialize_default,\n .deserialize = custom_deserialize_default,\n .compare_ext = custom_compare_ext_default,\n};\n\nstatic value alloc_surface(SDL_Surface *s)\n{\n value v = caml_alloc_custom(&surface_ops, sizeof(SDL_Surface *), 0, 1);\n Surface_val(v) = s;\n return v;\n}\n\n/* Texture */\n\n#define Texture_val(v) (*(SDL_Texture **)Data_custom_val(v))\n\nstatic void finalize_texture(value v)\n{\n SDL_Texture *t = Texture_val(v);\n if (t != NULL) { SDL_DestroyTexture(t); Texture_val(v) = NULL; }\n}\n\nstatic struct custom_operations texture_ops = {\n .identifier = \"usdl.texture\",\n .finalize = finalize_texture,\n .compare = custom_compare_default,\n .hash = custom_hash_default,\n .serialize = custom_serialize_default,\n .deserialize = custom_deserialize_default,\n .compare_ext = custom_compare_ext_default,\n};\n\nstatic value alloc_texture(SDL_Texture *t)\n{\n value v = caml_alloc_custom(&texture_ops, sizeof(SDL_Texture *), 0, 1);\n Texture_val(v) = t;\n return v;\n}\n\n/* Event — stored inline in custom block, no heap allocation */\n\n#define Event_val(v) ((SDL_Event *)Data_custom_val(v))\n\nstatic struct custom_operations event_ops = {\n .identifier = \"usdl.event\",\n .finalize = custom_finalize_default,\n .compare = custom_compare_default,\n .hash = custom_hash_default,\n .serialize = custom_serialize_default,\n .deserialize = custom_deserialize_default,\n .compare_ext = custom_compare_ext_default,\n};\n\n/* Init / quit */\n\nCAMLprim value caml_usdl_init(value vunit)\n{\n CAMLparam1(vunit);\n if (SDL_Init(SDL_INIT_VIDEO) < 0)\n caml_failwith(SDL_GetError());\n SDL_SetHint(SDL_HINT_RENDER_SCALE_QUALITY, \"1\");\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_usdl_quit(value vunit)\n{\n CAMLparam1(vunit);\n SDL_Quit();\n CAMLreturn(Val_unit);\n}\n\n/* Window */\n\nCAMLprim value caml_usdl_window_create(value vtitle, value vw, value vh)\n{\n CAMLparam3(vtitle, vw, vh);\n Uint32 flags = SDL_WINDOW_SHOWN | SDL_WINDOW_RESIZABLE |\n SDL_WINDOW_ALLOW_HIGHDPI;\n SDL_Window *win = SDL_CreateWindow(\n String_val(vtitle),\n SDL_WINDOWPOS_UNDEFINED, SDL_WINDOWPOS_UNDEFINED,\n Int_val(vw), Int_val(vh), flags);\n if (win == NULL) caml_failwith(SDL_GetError());\n CAMLreturn(alloc_window(win));\n}\n\nCAMLprim value caml_usdl_window_destroy(value vwin)\n{\n CAMLparam1(vwin);\n SDL_Window *w = Window_val(vwin);\n if (w != NULL) { SDL_DestroyWindow(w); Window_val(vwin) = NULL; }\n CAMLreturn(Val_unit);\n}\n\n/* Renderer */\n\nCAMLprim value caml_usdl_renderer_create(value vwin)\n{\n CAMLparam1(vwin);\n SDL_Window *win = Window_val(vwin);\n if (win == NULL) caml_invalid_argument(\"Usdl.Renderer.create: destroyed window\");\n Uint32 flags = SDL_RENDERER_ACCELERATED | SDL_RENDERER_PRESENTVSYNC;\n SDL_Renderer *ren = SDL_CreateRenderer(win, -1, flags);\n if (ren == NULL) caml_failwith(SDL_GetError());\n CAMLreturn(alloc_renderer(ren));\n}\n\nCAMLprim value caml_usdl_renderer_output_size(value vren)\n{\n CAMLparam1(vren);\n CAMLlocal1(result);\n SDL_Renderer *ren = Renderer_val(vren);\n if (ren == NULL) caml_invalid_argument(\"Usdl.Renderer.output_size: destroyed renderer\");\n int w, h;\n if (SDL_GetRendererOutputSize(ren, &w, &h) < 0)\n caml_failwith(SDL_GetError());\n result = caml_alloc_tuple(2);\n Store_field(result, 0, Val_int(w));\n Store_field(result, 1, Val_int(h));\n CAMLreturn(result);\n}\n\nCAMLprim value caml_usdl_renderer_clear(value vren)\n{\n CAMLparam1(vren);\n SDL_Renderer *ren = Renderer_val(vren);\n if (ren == NULL) caml_invalid_argument(\"Usdl.Renderer.clear: destroyed renderer\");\n if (SDL_RenderClear(ren) < 0) caml_failwith(SDL_GetError());\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_usdl_renderer_copy(value vren, value vtex)\n{\n CAMLparam2(vren, vtex);\n SDL_Renderer *ren = Renderer_val(vren);\n SDL_Texture *tex = Texture_val(vtex);\n if (ren == NULL || tex == NULL)\n caml_invalid_argument(\"Usdl.Renderer.copy: destroyed handle\");\n if (SDL_RenderCopy(ren, tex, NULL, NULL) < 0)\n caml_failwith(SDL_GetError());\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_usdl_renderer_present(value vren)\n{\n CAMLparam1(vren);\n SDL_Renderer *ren = Renderer_val(vren);\n if (ren == NULL) caml_invalid_argument(\"Usdl.Renderer.present: destroyed renderer\");\n SDL_RenderPresent(ren);\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_usdl_renderer_destroy(value vren)\n{\n CAMLparam1(vren);\n SDL_Renderer *r = Renderer_val(vren);\n if (r != NULL) { SDL_DestroyRenderer(r); Renderer_val(vren) = NULL; }\n CAMLreturn(Val_unit);\n}\n\n/* Surface */\n\nCAMLprim value caml_usdl_surface_create_argb8888(value vw, value vh)\n{\n CAMLparam2(vw, vh);\n SDL_Surface *s = SDL_CreateRGBSurfaceWithFormat(\n 0, Int_val(vw), Int_val(vh), 32, SDL_PIXELFORMAT_ARGB8888);\n if (s == NULL) caml_failwith(SDL_GetError());\n CAMLreturn(alloc_surface(s));\n}\n\nCAMLprim value caml_usdl_surface_destroy(value vsurf)\n{\n CAMLparam1(vsurf);\n SDL_Surface *s = Surface_val(vsurf);\n if (s != NULL) { SDL_FreeSurface(s); Surface_val(vsurf) = NULL; }\n CAMLreturn(Val_unit);\n}\n\nCAMLprim value caml_usdl_surface_pitch(value vsurf)\n{\n CAMLparam1(vsurf);\n SDL_Surface *s = Surface_val(vsurf);\n if (s == NULL) caml_invalid_argument(\"Usdl.Surface.pitch: destroyed surface\");\n CAMLreturn(Val_int(s->pitch));\n}\n\nCAMLprim value caml_usdl_surface_pixels(value vsurf)\n{\n CAMLparam1(vsurf);\n SDL_Surface *s = Surface_val(vsurf);\n if (s == NULL) caml_invalid_argument(\"Usdl.Surface.pixels: destroyed surface\");\n if (s->pixels == NULL) caml_failwith(\"Usdl.Surface.pixels: NULL pixels\");\n CAMLreturn(caml_ba_alloc_dims(\n CAML_BA_UINT8 | CAML_BA_C_LAYOUT | CAML_BA_EXTERNAL,\n 1, s->pixels, (intnat)s->h * s->pitch));\n}\n\n/* Texture */\n\nCAMLprim value caml_usdl_texture_of_surface(value vren, value vsurf)\n{\n CAMLparam2(vren, vsurf);\n SDL_Renderer *ren = Renderer_val(vren);\n SDL_Surface *s = Surface_val(vsurf);\n if (ren == NULL || s == NULL)\n caml_invalid_argument(\"Usdl.Texture.of_surface: destroyed handle\");\n SDL_Texture *tex = SDL_CreateTextureFromSurface(ren, s);\n if (tex == NULL) caml_failwith(SDL_GetError());\n CAMLreturn(alloc_texture(tex));\n}\n\nCAMLprim value caml_usdl_texture_destroy(value vtex)\n{\n CAMLparam1(vtex);\n SDL_Texture *t = Texture_val(vtex);\n if (t != NULL) { SDL_DestroyTexture(t); Texture_val(vtex) = NULL; }\n CAMLreturn(Val_unit);\n}\n\n/* Event */\n\nCAMLprim value caml_usdl_event_create(value vunit)\n{\n CAMLparam1(vunit);\n value v = caml_alloc_custom(&event_ops, sizeof(SDL_Event), 0, 1);\n memset(Event_val(v), 0, sizeof(SDL_Event));\n CAMLreturn(v);\n}\n\nCAMLprim value caml_usdl_event_wait(value vev)\n{\n CAMLparam1(vev);\n SDL_Event ev;\n caml_release_runtime_system();\n int ret = SDL_WaitEvent(&ev);\n caml_acquire_runtime_system();\n if (ret == 1) memcpy(Event_val(vev), &ev, sizeof(SDL_Event));\n CAMLreturn(Val_bool(ret == 1));\n}\n\nCAMLprim value caml_usdl_event_type(value vev)\n{\n return Val_int(Event_val(vev)->type);\n}\n\nCAMLprim value caml_usdl_event_window_id(value vev)\n{\n SDL_Event *ev = Event_val(vev);\n if (ev->type == SDL_WINDOWEVENT)\n return Val_int(ev->window.event);\n return Val_int(-1);\n}\n\nCAMLprim value caml_usdl_event_keycode(value vev)\n{\n SDL_Event *ev = Event_val(vev);\n if (ev->type == SDL_KEYDOWN || ev->type == SDL_KEYUP)\n return Val_int(ev->key.keysym.sym);\n return Val_int(-1);\n}\n"
},
{
"path": "packages/kaun/README.md",
"content": "# Kaun\n\nNeural networks and training utilities for OCaml, built on [Rune](../rune/)\n\nKaun provides composable layers, optimizers with learning-rate schedules,\nautomatic differentiation over parameter trees, data pipelines, and a\nhigh-level training loop. It also supports loading pretrained models from\nHuggingFace.\n\n## Quick Start\n\nTrain a small network on XOR:\n\n```ocaml\nopen Kaun\n\nlet () =\n let rngs = Rune.Rng.key 42 in\n let dtype = Rune.float32 in\n\n let x = Rune.create dtype [| 4; 2 |] [| 0.; 0.; 0.; 1.; 1.; 0.; 1.; 1. |] in\n let y = Rune.create dtype [| 4; 1 |] [| 0.; 1.; 1.; 0. |] in\n\n let model =\n Layer.sequential\n [\n Layer.linear ~in_features:2 ~out_features:4 ();\n Layer.tanh ();\n Layer.linear ~in_features:4 ~out_features:1 ();\n ]\n in\n let trainer =\n Train.make ~model\n ~optimizer:(Optim.adam ~lr:(Optim.Schedule.constant 0.01) ())\n in\n let st = Train.init trainer ~rngs ~dtype in\n let st =\n Train.fit trainer st ~rngs\n ~report:(fun ~step ~loss _st ->\n if step mod 200 = 0 then Printf.printf \"step %4d loss %.6f\\n\" step loss)\n (Data.repeat 1000 (x, fun pred -> Loss.binary_cross_entropy pred y))\n in\n let pred = Train.predict trainer st x |> Rune.sigmoid in\n for i = 0 to 3 do\n Printf.printf \" [%.0f, %.0f] -> %.3f\\n\"\n (Rune.item [ i; 0 ] x) (Rune.item [ i; 1 ] x) (Rune.item [ i; 0 ] pred)\n done\n```\n\n## Features\n\n- **Layers**: linear, conv1d, conv2d, layer norm, RMS norm, batch norm, embedding, dropout, multi-head attention with RoPE, and all standard activations (relu, gelu, tanh, sigmoid, etc.)\n- **Composition**: `Layer.sequential` and `Layer.compose` for building models\n- **Optimizers**: SGD, Adam, AdamW, RMSprop, Adagrad with gradient clipping\n- **Schedules**: constant, cosine decay, warmup cosine, exponential decay, warmup linear\n- **Training**: `Train.fit` iterates over `Data.t` pipelines with early stopping and per-step reporting; `Train.step` for manual control\n- **Data pipelines**: lazy, composable iterators with shuffle, batching, and `Data.prepare` for the common (x, y) tensor pair workflow\n- **Metrics**: running trackers, dataset evaluation, accuracy, precision, recall, F1\n- **Losses**: cross-entropy, sparse cross-entropy, binary cross-entropy, MSE, MAE\n- **Parameter trees**: `Ptree.t` for heterogeneous tensor storage, mapping, and serialization\n- **Checkpointing**: save/load to SafeTensors format\n- **HuggingFace**: download pretrained weights and configs (`kaun.hf`)\n- **Datasets**: MNIST and FashionMNIST loaders (`kaun.datasets`)\n\n## Libraries\n\n| Library | opam package | Description |\n|---------|-------------|-------------|\n| `kaun` | `kaun` | Core: layers, optimizers, training, data, metrics |\n| `kaun_hf` | `kaun.hf` | HuggingFace Hub integration |\n| `kaun_datasets` | `kaun.datasets` | Dataset loaders (MNIST, FashionMNIST) |\n\n## Examples\n\n- **01-xor** -- Binary classification on XOR with a 2-layer network\n- **02-mnist** -- CNN with conv2d, pooling, and multi-epoch training on MNIST\n- **03-bert** -- Fine-tune pretrained BERT for sentiment classification\n- **04-gpt2** -- Autoregressive text generation with pretrained GPT-2\n\n## Contributing\n\nSee the [Raven monorepo README](../README.md) for guidelines.\n\n## License\n\nISC License. See [LICENSE](../LICENSE) for details.\n"
},
{
"path": "packages/kaun/bench/bench_kaun.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Fn = Kaun.Fn\nmodule Layer = Kaun.Layer\nmodule Loss = Kaun.Loss\nmodule Attention = Kaun.Attention\n\nlet batch = 32\nlet seq_len = 128\nlet dim = 256\nlet num_heads = 8\nlet num_classes = 10\n\nlet normalization_benchmarks () =\n let x = Nx.rand Nx.float32 [| batch; seq_len; dim |] in\n let gamma = Nx.ones Nx.float32 [| dim |] in\n let beta = Nx.zeros Nx.float32 [| dim |] in\n let x_4d = Nx.rand Nx.float32 [| batch; dim; 8; 8 |] in\n let bn4_scale = Nx.ones Nx.float32 [| dim |] in\n let bn4_bias = Nx.zeros Nx.float32 [| dim |] in\n let x_2d = Nx.rand Nx.float32 [| batch; dim |] in\n let bn2_scale = Nx.ones Nx.float32 [| dim |] in\n let bn2_bias = Nx.zeros Nx.float32 [| dim |] in\n [\n Thumper.bench \"layer_norm [32;128;256]\" (fun () ->\n Fn.layer_norm ~gamma ~beta x);\n Thumper.bench \"rms_norm [32;128;256]\" (fun () -> Fn.rms_norm x);\n Thumper.bench \"batch_norm [32;256;8;8]\" (fun () ->\n Fn.batch_norm ~scale:bn4_scale ~bias:bn4_bias x_4d);\n Thumper.bench \"batch_norm [32;256]\" (fun () ->\n Fn.batch_norm ~scale:bn2_scale ~bias:bn2_bias x_2d);\n ]\n\nlet attention_benchmarks () =\n let head_dim = dim / num_heads in\n let q = Nx.rand Nx.float32 [| batch; num_heads; seq_len; head_dim |] in\n let k = Nx.rand Nx.float32 [| batch; num_heads; seq_len; head_dim |] in\n let v = Nx.rand Nx.float32 [| batch; num_heads; seq_len; head_dim |] in\n let x = Nx.rand Nx.float32 [| batch; seq_len; head_dim |] in\n [\n Thumper.bench \"dot_product_attention [32;8;128;32]\" (fun () ->\n Fn.dot_product_attention q k v);\n Thumper.bench \"dot_product_attention causal [32;8;128;32]\" (fun () ->\n Fn.dot_product_attention ~is_causal:true q k v);\n Thumper.bench \"rope [32;128;32]\" (fun () -> Attention.rope x);\n ]\n\nlet loss_benchmarks () =\n let logits = Nx.rand Nx.float32 [| batch; num_classes |] in\n let labels_onehot = Nx.zeros Nx.float32 [| batch; num_classes |] in\n let targets = Nx.rand Nx.float32 [| batch; num_classes |] in\n let predictions = Nx.rand Nx.float32 [| batch; num_classes |] in\n let binary_logits = Nx.rand Nx.float32 [| batch |] in\n let binary_labels = Nx.zeros Nx.float32 [| batch |] in\n [\n Thumper.bench \"cross_entropy [32;10]\" (fun () ->\n Loss.cross_entropy logits labels_onehot);\n Thumper.bench \"binary_cross_entropy [32]\" (fun () ->\n Loss.binary_cross_entropy binary_logits binary_labels);\n Thumper.bench \"mse [32;10]\" (fun () -> Loss.mse predictions targets);\n Thumper.bench \"mae [32;10]\" (fun () -> Loss.mae predictions targets);\n ]\n\nlet conv_benchmarks () =\n let x1d = Nx.rand Nx.float32 [| batch; 64; 128 |] in\n let w1d = Nx.rand Nx.float32 [| 128; 64; 3 |] in\n let x2d = Nx.rand Nx.float32 [| batch; 64; 32; 32 |] in\n let w2d = Nx.rand Nx.float32 [| 128; 64; 3; 3 |] in\n [\n Thumper.bench \"conv1d [32;64;128] k=3\" (fun () -> Fn.conv1d x1d w1d);\n Thumper.bench \"conv1d same [32;64;128] k=3\" (fun () ->\n Fn.conv1d ~padding:`Same x1d w1d);\n Thumper.bench \"conv2d [32;64;32;32] k=3x3\" (fun () -> Fn.conv2d x2d w2d);\n Thumper.bench \"conv2d same [32;64;32;32] k=3x3\" (fun () ->\n Fn.conv2d ~padding:`Same x2d w2d);\n ]\n\nlet pooling_benchmarks () =\n let x2d = Nx.rand Nx.float32 [| batch; 64; 32; 32 |] in\n [\n Thumper.bench \"max_pool2d [32;64;32;32] k=2x2\" (fun () ->\n Fn.max_pool2d ~kernel_size:(2, 2) x2d);\n Thumper.bench \"avg_pool2d [32;64;32;32] k=2x2\" (fun () ->\n Fn.avg_pool2d ~kernel_size:(2, 2) x2d);\n ]\n\nlet layer_benchmarks () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let linear_layer = Layer.linear ~in_features:dim ~out_features:dim () in\n let linear_vars = Layer.init linear_layer ~dtype:Nx.float32 in\n let ln_layer = Layer.layer_norm ~dim () in\n let ln_vars = Layer.init ln_layer ~dtype:Nx.float32 in\n let mha_layer = Attention.multi_head_attention ~embed_dim:dim ~num_heads () in\n let mha_vars = Layer.init mha_layer ~dtype:Nx.float32 in\n let x = Nx.rand Nx.float32 [| batch; seq_len; dim |] in\n [\n Thumper.bench \"Layer.linear [32;128;256]->[32;128;256]\" (fun () ->\n Layer.apply linear_layer linear_vars ~training:false x);\n Thumper.bench \"Layer.layer_norm [32;128;256]\" (fun () ->\n Layer.apply ln_layer ln_vars ~training:false x);\n Thumper.bench \"Layer.multi_head_attention [32;128;256] h=8\" (fun () ->\n Layer.apply mha_layer mha_vars ~training:false x);\n ]\n\nlet embedding_benchmarks () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let vocab_size = 32000 in\n let embed_dim = dim in\n let table = Nx.rand Nx.float32 [| vocab_size; embed_dim |] in\n let indices =\n Nx.create Nx.int32 [| batch; seq_len |]\n (Array.init (batch * seq_len) (fun i -> Int32.of_int (i mod vocab_size)))\n in\n [\n Thumper.bench \"embedding [32;128] vocab=32000 dim=256\" (fun () ->\n Fn.embedding ~embedding:table indices);\n ]\n\nlet build_benchmarks () =\n [\n Thumper.group \"Normalization\" (normalization_benchmarks ());\n Thumper.group \"Attention\" (attention_benchmarks ());\n Thumper.group \"Loss\" (loss_benchmarks ());\n Thumper.group \"Convolution\" (conv_benchmarks ());\n Thumper.group \"Pooling\" (pooling_benchmarks ());\n Thumper.group \"Layer\" (layer_benchmarks ());\n Thumper.group \"Embedding\" (embedding_benchmarks ());\n ]\n\nlet () =\n let benchmarks = build_benchmarks () in\n Thumper.run \"kaun\" benchmarks\n"
},
{
"path": "packages/kaun/bench/dune",
"content": "(executable\n (name bench_kaun)\n (libraries nx rune kaun thumper))\n\n(rule\n (alias runtest)\n (action\n (progn\n (run %{exe:bench_kaun.exe} -q)\n (diff? kaun.thumper kaun.thumper.corrected))))\n"
},
{
"path": "packages/kaun/bench/kaun.thumper",
"content": "# thumper baseline\n# version: 1\n# suite_name: kaun\n# host: 1480401c3b76ed18\n# cpu: Apple M1 Max\n# ocaml: 5.4.1\n# git: 31747323\n# dirty: true\n# command: /Users/tmattio/Workspace/raven/_build/default/packages/kaun/bench/bench_kaun.exe --bless --quick\n\nattention/dot_product_attention__32_8_128_32_\talloc_words\t2.137000e+03\t2.137000e+03\t2.137000e+03\t0.000000e+00\t5\t1\nattention/dot_product_attention__32_8_128_32_\tcpu_time\t9.671033e-02\t9.639139e-02\t9.768167e-02\t6.670834e-03\t5\t0\nattention/dot_product_attention__32_8_128_32_\twall_time\t7.475045e-02\t7.419780e-02\t7.636849e-02\t1.451961e-02\t5\t0\nattention/dot_product_attention_causal__32_8_128_32_\talloc_words\t8.969000e+03\t8.969000e+03\t8.969000e+03\t0.000000e+00\t5\t1\nattention/dot_product_attention_causal__32_8_128_32_\tcpu_time\t1.370568e-01\t1.362397e-01\t1.373701e-01\t4.123616e-03\t5\t0\nattention/dot_product_attention_causal__32_8_128_32_\twall_time\t1.113716e-01\t1.109278e-01\t1.116209e-01\t3.111817e-03\t5\t0\nattention/rope__32_128_32_\talloc_words\t6.995000e+03\t6.995000e+03\t6.995000e+03\t0.000000e+00\t5\t0\nattention/rope__32_128_32_\tcpu_time\t4.598942e-03\t4.544050e-03\t4.656300e-03\t1.220391e-02\t5\t0\nattention/rope__32_128_32_\twall_time\t3.115630e-03\t3.091075e-03\t3.141427e-03\t8.080598e-03\t5\t0\nconvolution/conv1d__32_64_128__k_3\talloc_words\t5.000000e+02\t5.000000e+02\t5.000000e+02\t0.000000e+00\t5\t0\nconvolution/conv1d__32_64_128__k_3\tcpu_time\t1.620888e-03\t1.596007e-03\t1.637217e-03\t1.271213e-02\t5\t0\nconvolution/conv1d__32_64_128__k_3\twall_time\t1.626572e-03\t1.602353e-03\t1.644431e-03\t1.293450e-02\t5\t0\nconvolution/conv1d_same__32_64_128__k_3\talloc_words\t5.100000e+02\t5.100000e+02\t5.100000e+02\t0.000000e+00\t5\t0\nconvolution/conv1d_same__32_64_128__k_3\tcpu_time\t5.615602e-03\t5.517316e-03\t5.660000e-03\t1.270422e-02\t5\t2\nconvolution/conv1d_same__32_64_128__k_3\twall_time\t1.344058e-03\t1.329492e-03\t1.358392e-03\t1.075103e-02\t5\t0\nconvolution/conv2d__32_64_32_32__k_3x3\talloc_words\t6.070000e+02\t6.070000e+02\t6.070000e+02\t0.000000e+00\t5\t1\nconvolution/conv2d__32_64_32_32__k_3x3\tcpu_time\t2.908078e-02\t2.856329e-02\t2.982455e-02\t2.168548e-02\t5\t0\nconvolution/conv2d__32_64_32_32__k_3x3\twall_time\t3.174315e-02\t2.900269e-02\t3.589196e-02\t1.085158e-01\t5\t1\nconvolution/conv2d_same__32_64_32_32__k_3x3\talloc_words\t6.200000e+02\t6.200000e+02\t6.200000e+02\t0.000000e+00\t5\t1\nconvolution/conv2d_same__32_64_32_32__k_3x3\tcpu_time\t1.506193e-01\t1.496805e-01\t1.521535e-01\t8.209534e-03\t5\t0\nconvolution/conv2d_same__32_64_32_32__k_3x3\twall_time\t3.016738e-02\t2.915340e-02\t3.056394e-02\t2.337867e-02\t5\t1\nembedding/embedding__32_128__vocab_32000_dim_256\talloc_words\t9.950000e+02\t9.950000e+02\t9.950000e+02\t0.000000e+00\t5\t0\nembedding/embedding__32_128__vocab_32000_dim_256\tcpu_time\t7.132733e-03\t7.093132e-03\t7.245324e-03\t1.066859e-02\t5\t1\nembedding/embedding__32_128__vocab_32000_dim_256\twall_time\t6.695677e-03\t6.659470e-03\t6.761231e-03\t7.598951e-03\t5\t1\nlayer/layer_layer_norm__32_128_256_\talloc_words\t3.773000e+03\t3.773000e+03\t3.773000e+03\t0.000000e+00\t5\t1\nlayer/layer_layer_norm__32_128_256_\tcpu_time\t2.918713e-02\t2.908128e-02\t2.951400e-02\t7.412812e-03\t5\t1\nlayer/layer_layer_norm__32_128_256_\twall_time\t2.688614e-02\t2.676829e-02\t2.714231e-02\t6.955731e-03\t5\t2\nlayer/layer_linear__32_128_256_-__32_128_256_\talloc_words\t5.640000e+02\t5.640000e+02\t5.640000e+02\t0.000000e+00\t5\t0\nlayer/layer_linear__32_128_256_-__32_128_256_\tcpu_time\t7.458522e-03\t7.400407e-03\t7.510652e-03\t7.390491e-03\t5\t0\nlayer/layer_linear__32_128_256_-__32_128_256_\twall_time\t6.948420e-03\t6.920065e-03\t6.969745e-03\t3.574904e-03\t5\t0\nlayer/layer_multi_head_attention__32_128_256__h_8\talloc_words\t3.827000e+03\t3.827000e+03\t3.827000e+03\t0.000000e+00\t5\t1\nlayer/layer_multi_head_attention__32_128_256__h_8\tcpu_time\t1.053272e-01\t1.052223e-01\t1.058294e-01\t2.882088e-03\t5\t0\nlayer/layer_multi_head_attention__32_128_256__h_8\twall_time\t8.404089e-02\t8.353135e-02\t8.467066e-02\t6.778344e-03\t5\t1\nloss/binary_cross_entropy__32_\talloc_words\t5.785000e+03\t5.785000e+03\t5.785000e+03\t0.000000e+00\t5\t0\nloss/binary_cross_entropy__32_\tcpu_time\t2.444898e-05\t2.428635e-05\t2.461726e-05\t6.767379e-03\t5\t0\nloss/binary_cross_entropy__32_\twall_time\t2.458928e-05\t2.440828e-05\t2.479800e-05\t7.924565e-03\t5\t0\nloss/cross_entropy__32_10_\talloc_words\t1.968000e+03\t1.968000e+03\t1.968000e+03\t0.000000e+00\t5\t0\nloss/cross_entropy__32_10_\tcpu_time\t1.317014e-05\t1.291714e-05\t1.357537e-05\t2.498951e-02\t5\t1\nloss/cross_entropy__32_10_\twall_time\t1.346084e-05\t1.302917e-05\t1.403663e-05\t3.742195e-02\t5\t1\nloss/mae__32_10_\talloc_words\t6.390000e+02\t6.390000e+02\t6.390000e+02\t0.000000e+00\t6\t0\nloss/mae__32_10_\tcpu_time\t3.343846e-06\t3.255477e-06\t3.435480e-06\t2.691551e-02\t6\t0\nloss/mae__32_10_\twall_time\t3.398029e-06\t3.290374e-06\t3.521626e-06\t3.402740e-02\t6\t0\nloss/mse__32_10_\talloc_words\t7.030000e+02\t7.030000e+02\t7.030000e+02\t0.000000e+00\t5\t0\nloss/mse__32_10_\tcpu_time\t3.414987e-06\t3.392142e-06\t3.434869e-06\t6.255854e-03\t5\t0\nloss/mse__32_10_\twall_time\t3.428576e-06\t3.401322e-06\t3.456842e-06\t8.096629e-03\t5\t0\nnormalization/batch_norm__32_256_\talloc_words\t4.656000e+03\t4.656000e+03\t4.656000e+03\t0.000000e+00\t5\t0\nnormalization/batch_norm__32_256_\tcpu_time\t2.447512e-04\t2.412132e-04\t2.483533e-04\t1.458635e-02\t5\t0\nnormalization/batch_norm__32_256_\twall_time\t2.451437e-04\t2.417406e-04\t2.493240e-04\t1.546732e-02\t5\t0\nnormalization/batch_norm__32_256_8_8_\talloc_words\t5.006000e+03\t5.006000e+03\t5.006000e+03\t0.000000e+00\t10\t1\nnormalization/batch_norm__32_256_8_8_\tcpu_time\t2.507420e-02\t2.377102e-02\t2.591671e-02\t4.278692e-02\t10\t0\nnormalization/batch_norm__32_256_8_8_\twall_time\t2.386942e-02\t2.278693e-02\t2.460672e-02\t3.811965e-02\t10\t0\nnormalization/layer_norm__32_128_256_\talloc_words\t3.754000e+03\t3.754000e+03\t3.754000e+03\t0.000000e+00\t5\t1\nnormalization/layer_norm__32_128_256_\tcpu_time\t2.951118e-02\t2.887281e-02\t2.970830e-02\t1.415546e-02\t5\t0\nnormalization/layer_norm__32_128_256_\twall_time\t2.723408e-02\t2.674500e-02\t2.737987e-02\t1.165582e-02\t5\t0\nnormalization/rms_norm__32_128_256_\talloc_words\t1.868000e+03\t1.868000e+03\t1.868000e+03\t0.000000e+00\t5\t0\nnormalization/rms_norm__32_128_256_\tcpu_time\t7.965700e-03\t7.931148e-03\t7.994719e-03\t3.990293e-03\t5\t0\nnormalization/rms_norm__32_128_256_\twall_time\t7.253412e-03\t7.224232e-03\t7.272761e-03\t3.345277e-03\t5\t2\npooling/avg_pool2d__32_64_32_32__k_2x2\talloc_words\t9.650000e+02\t9.650000e+02\t9.650000e+02\t0.000000e+00\t5\t1\npooling/avg_pool2d__32_64_32_32__k_2x2\tcpu_time\t2.944396e-02\t2.936873e-02\t2.954386e-02\t2.973982e-03\t5\t0\npooling/avg_pool2d__32_64_32_32__k_2x2\twall_time\t2.400500e-02\t2.392232e-02\t2.405630e-02\t2.790660e-03\t5\t0\npooling/max_pool2d__32_64_32_32__k_2x2\talloc_words\t4.980000e+02\t4.980000e+02\t4.980000e+02\t0.000000e+00\t5\t1\npooling/max_pool2d__32_64_32_32__k_2x2\tcpu_time\t7.406087e-02\t7.315733e-02\t7.438250e-02\t8.271349e-03\t5\t0\npooling/max_pool2d__32_64_32_32__k_2x2\twall_time\t2.149041e-02\t2.079528e-02\t2.253364e-02\t4.044505e-02\t5\t1\n"
},
{
"path": "packages/kaun/doc/01-getting-started.md",
"content": "# Getting Started\n\nThis guide covers installation, key concepts, and two complete examples:\nlearning XOR and classifying MNIST digits.\n\n## Installation\n\n\n```bash\nopam install kaun\n```\n\nOr build from source:\n\n\n```bash\ngit clone https://github.com/raven-ml/raven\ncd raven && dune build kaun\n```\n\n## Key Concepts\n\n**Layer.** A layer is a record `{ init; apply }`. `init` creates fresh\nparameters and state. `apply` runs the forward pass. Layers compose with\n`Layer.sequential` (homogeneous float pipelines) and `Layer.compose`\n(heterogeneous, e.g. embedding to dense).\n\n**Ptree.** A `Ptree.t` is a tree of tensors. Dict nodes hold named\nsubtrees, list nodes hold ordered subtrees, and leaves hold tensors.\nParameters and state are both `Ptree.t` values — plain data you can\ninspect, map, serialize, and load.\n\n**Layer.vars.** A `vars` bundles `params` (trainable), `state`\n(non-trainable, e.g. batch norm running statistics), and a `dtype`\nwitness.\n\n**Train.** `Train.make` pairs a model with an optimizer. `Train.init`\ncreates the initial training state. `Train.fit` trains over a `Data.t`\npipeline. `Train.predict` runs inference.\n\n**Data.** `Data.t` is a lazy, composable iterator. Build from tensors or\narrays, shuffle, batch, map, and feed to `Train.fit`.\n\n**Optim.** An optimizer combines a learning-rate schedule with an update\nrule. Schedules are functions `int -> float`.\n\n## Example: XOR\n\nThe XOR problem is the simplest non-linear classification task. This\nexample trains a small network to learn it.\n\n\n```ocaml\nopen Kaun\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n\n (* XOR dataset: 4 examples, 2 features each *)\n let x = Nx.create Nx.Float32 [| 4; 2 |] [| 0.; 0.; 0.; 1.; 1.; 0.; 1.; 1. |] in\n let y = Nx.create Nx.Float32 [| 4; 1 |] [| 0.; 1.; 1.; 0. |] in\n\n (* Model: 2 -> 4 -> 1 with tanh activation *)\n let model =\n Layer.sequential\n [\n Layer.linear ~in_features:2 ~out_features:4 ();\n Layer.tanh ();\n Layer.linear ~in_features:4 ~out_features:1 ();\n ]\n in\n\n (* Create a trainer: model + optimizer *)\n let trainer =\n Train.make ~model\n ~optimizer:(Vega.adam (Vega.Schedule.constant 0.01))\n in\n\n (* Initialize training state (model vars + optimizer state) *)\n let st = Train.init trainer ~dtype:Nx.Float32 in\n\n (* Train for 1000 steps on the same data *)\n let st =\n Train.fit trainer st\n ~report:(fun ~step ~loss _st ->\n if step mod 200 = 0 then\n Printf.printf \"step %4d loss %.6f\\n\" step loss)\n (Data.repeat 1000 (x, fun pred -> Loss.binary_cross_entropy pred y))\n in\n\n (* Predict *)\n let pred = Train.predict trainer st x |> Nx.sigmoid in\n Printf.printf \"\\npredictions (expected 0 1 1 0):\\n\";\n for i = 0 to 3 do\n Printf.printf \" [%.0f, %.0f] -> %.3f\\n\"\n (Nx.item [ i; 0 ] x)\n (Nx.item [ i; 1 ] x)\n (Nx.item [ i; 0 ] pred)\n done\n```\n\nKey points:\n\n- `Data.repeat 1000 (x, loss_fn)` creates a pipeline that yields the\n same `(input, loss_fn)` pair 1000 times.\n- The loss function `fun pred -> Loss.binary_cross_entropy pred y`\n receives the model output and computes a scalar loss.\n- `Train.predict` runs in evaluation mode (no dropout, no state\n updates).\n\n## Example: MNIST\n\nA convolutional network for handwritten digit classification using the\nbuilt-in MNIST dataset loader.\n\n\n```ocaml\nopen Kaun\n\nlet batch_size = 64\nlet epochs = 3\nlet lr = 0.001\n\nlet model =\n Layer.sequential\n [\n Layer.conv2d ~in_channels:1 ~out_channels:16 ();\n Layer.relu ();\n Layer.max_pool2d ~kernel_size:(2, 2) ();\n Layer.conv2d ~in_channels:16 ~out_channels:32 ();\n Layer.relu ();\n Layer.max_pool2d ~kernel_size:(2, 2) ();\n Layer.flatten ();\n Layer.linear ~in_features:(32 * 7 * 7) ~out_features:128 ();\n Layer.relu ();\n Layer.linear ~in_features:128 ~out_features:10 ();\n ]\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n\n Printf.printf \"Loading MNIST...\\n%!\";\n let (x_train, y_train), (x_test, y_test) = Kaun_datasets.mnist () in\n let n_train = (Nx.shape x_train).(0) in\n Printf.printf \" train: %d test: %d\\n%!\" n_train (Nx.shape x_test).(0);\n\n (* Fixed test batches *)\n let test_batches = Data.prepare ~batch_size (x_test, y_test) in\n\n (* Trainer *)\n let trainer =\n Train.make ~model\n ~optimizer:(Vega.adam (Vega.Schedule.constant lr))\n in\n let st = ref (Train.init trainer ~dtype:Nx.Float32) in\n\n for epoch = 1 to epochs do\n (* Shuffle training data each epoch *)\n let train_data =\n Data.prepare ~shuffle:true ~batch_size (x_train, y_train)\n |> Data.map (fun (x, y) ->\n (x, fun logits -> Loss.cross_entropy_sparse logits y))\n in\n let num_batches = n_train / batch_size in\n let tracker = Metric.tracker () in\n\n st :=\n Train.fit trainer !st\n ~report:(fun ~step ~loss _st ->\n Metric.observe tracker \"loss\" loss;\n Printf.printf \"\\r batch %d/%d loss: %.4f%!\" step num_batches loss)\n train_data;\n Printf.printf \"\\n%!\";\n\n (* Evaluate on test set *)\n Data.reset test_batches;\n let test_acc =\n Metric.eval\n (fun (x, y) ->\n let logits = Train.predict trainer !st x in\n Metric.accuracy logits y)\n test_batches\n in\n Printf.printf \"epoch %d train_loss: %.4f test_acc: %.2f%%\\n%!\" epoch\n (Metric.mean tracker \"loss\")\n (test_acc *. 100.)\n done\n```\n\nKey points:\n\n- `Kaun_datasets.mnist ()` returns `((x_train, y_train), (x_test, y_test))`\n as float32 tensor pairs. Images have shape `[N; 1; 28; 28]` (NCHW),\n labels `[N]`.\n- `Data.prepare ~shuffle:key ~batch_size (x, y)` creates a shuffled,\n batched pipeline of tensor pairs.\n- `Data.map` attaches the loss function to each batch, producing the\n `(input, loss_fn)` pairs that `Train.fit` expects.\n- `Metric.eval` folds a function over a data pipeline and returns the\n mean.\n- `Metric.tracker` accumulates running means for reporting.\n\n## Next Steps\n\n- [Layers and Models](../02-layers-and-models/) — full layer catalog, composition patterns, custom layers\n- [Training](../03-training/) — optimizers, schedules, losses, data pipelines, custom loops\n- [Checkpoints and Pretrained Models](../04-checkpoints-and-pretrained/) — saving, loading, HuggingFace Hub\n"
},
{
"path": "packages/kaun/doc/02-layers-and-models.md",
"content": "# Layers and Models\n\nA `Layer.t` pairs parameter initialization with a forward computation.\nThis guide covers the built-in layers, composition, custom layers, and\nthe `vars` type.\n\n## The Layer Type\n\nA layer is a record with two fields:\n\n\n```ocaml\ntype ('input, 'output) Layer.t = {\n init :\n 'layout.\n dtype:(float, 'layout) Nx.dtype -> 'layout vars;\n apply :\n 'layout 'in_elt.\n params:Ptree.t ->\n state:Ptree.t ->\n dtype:(float, 'layout) Nx.dtype ->\n training:bool ->\n ?ctx:Context.t ->\n ('input, 'in_elt) Nx.t ->\n ('output, 'layout) Nx.t * Ptree.t;\n}\n```\n\nThe type parameters `'input` and `'output` describe the element types.\nMost layers use `(float, float) Layer.t` — they accept and produce float\ntensors. `embedding` is `(int32, float) Layer.t` — it accepts int32\nindices and produces float vectors.\n\nUse `Layer.init` and `Layer.apply` instead of accessing fields directly:\n\n\n```ocaml\nlet vars = Layer.init model ~dtype:Nx.Float32\nlet output, vars' = Layer.apply model vars ~training:false x\n```\n\n## The vars Type\n\n`Layer.vars` bundles trainable parameters, non-trainable state, and a\ndtype witness:\n\n\n```ocaml\nLayer.params vars (* Ptree.t — trainable parameters *)\nLayer.state vars (* Ptree.t — non-trainable state (e.g. batch norm stats) *)\nLayer.dtype vars (* dtype witness *)\n```\n\nUse `Layer.with_params` and `Layer.with_state` to replace components:\n\n\n```ocaml\nlet vars' = Layer.with_params vars new_params\n```\n\n## Composition\n\n### sequential\n\n`Layer.sequential` chains `(float, float) Layer.t` layers in order.\nParameters are stored as a `Ptree.List`:\n\n\n```ocaml\nlet model = Layer.sequential [\n Layer.linear ~in_features:784 ~out_features:128 ();\n Layer.relu ();\n Layer.linear ~in_features:128 ~out_features:10 ();\n]\n```\n\n### compose\n\n`Layer.compose` chains two layers with different input/output types.\nParameters are stored as a `Ptree.Dict` with keys `\"left\"` and\n`\"right\"`:\n\n\n```ocaml\n(* embedding (int32 -> float) composed with a linear layer (float -> float) *)\nlet embed_then_project =\n Layer.compose\n (Layer.embedding ~vocab_size:10000 ~embed_dim:256 ())\n (Layer.linear ~in_features:256 ~out_features:128 ())\n(* embed_then_project : (int32, float) Layer.t *)\n```\n\n## Dense\n\n\n```ocaml\nLayer.linear ~in_features:784 ~out_features:128 ()\n```\n\nFully connected layer computing `xW + b`. Optional `~weight_init` and\n`~bias_init` arguments override the defaults (Glorot uniform for\nweights, zeros for bias).\n\n## Convolution\n\n\n```ocaml\n(* 1D: input [batch; in_channels; length] *)\nLayer.conv1d ~in_channels:3 ~out_channels:16 ()\nLayer.conv1d ~in_channels:3 ~out_channels:16 ~kernel_size:5 ~stride:2 ~padding:`Valid ()\n\n(* 2D: input [batch; in_channels; height; width] *)\nLayer.conv2d ~in_channels:1 ~out_channels:32 ()\nLayer.conv2d ~in_channels:1 ~out_channels:32 ~kernel_size:(5, 5) ()\n```\n\n`conv1d` supports configurable `~kernel_size` (default 3), `~stride`\n(default 1), `~dilation` (default 1), and `~padding` (default `` `Same ``).\n\n`conv2d` supports configurable `~kernel_size` (default `(3, 3)`). Stride\nis `(1, 1)` and padding is `` `Same ``.\n\n## Normalization\n\n\n```ocaml\nLayer.layer_norm ~dim:128 () (* learnable gamma and beta *)\nLayer.layer_norm ~dim:128 ~eps:1e-6 ()\n\nLayer.rms_norm ~dim:128 () (* learnable scale, no bias *)\n\nLayer.batch_norm ~num_features:32 () (* learnable scale and bias,\n running mean/var in state *)\n```\n\n`batch_norm` updates running statistics during training and uses them\nduring evaluation. Normalization axes are inferred from rank: rank 2\nuses `[0]`, rank 3 uses `[0; 2]`, rank 4 uses `[0; 2; 3]`.\n\n## Embedding\n\n\n```ocaml\nLayer.embedding ~vocab_size:10000 ~embed_dim:256 ()\n```\n\nInput: int32 token indices of any shape. Output: float tensors with\n`embed_dim` appended to the input shape.\n\nWhen `~scale:true` (the default), output vectors are multiplied by\n`sqrt(embed_dim)`.\n\n## Regularization\n\n\n```ocaml\nLayer.dropout ~rate:0.1 ()\n```\n\nDuring training (`~training:true`), randomly zeros elements with\nprobability `rate`. Requires `~rngs` during training. Identity during\nevaluation.\n\n## Activations\n\nAll activation layers have no parameters:\n\n\n```ocaml\nLayer.relu () (* max(0, x) *)\nLayer.gelu () (* Gaussian error linear unit *)\nLayer.silu () (* x * sigmoid(x) *)\nLayer.tanh () (* hyperbolic tangent *)\nLayer.sigmoid () (* logistic function *)\n```\n\n## Pooling\n\n\n```ocaml\nLayer.max_pool2d ~kernel_size:(2, 2) ()\nLayer.avg_pool2d ~kernel_size:(2, 2) ()\nLayer.max_pool2d ~kernel_size:(2, 2) ~stride:(1, 1) ()\n```\n\n`~stride` defaults to `~kernel_size`. No parameters.\n\n## Reshape\n\n\n```ocaml\nLayer.flatten ()\n```\n\nFlattens all dimensions after the batch dimension:\n`[batch; d1; ...; dn]` becomes `[batch; d1 * ... * dn]`.\n\n## Multi-Head Attention\n\n\n```ocaml\nAttention.multi_head_attention ~embed_dim:256 ~num_heads:8 ()\n```\n\nInput shape: `[batch; seq_len; embed_dim]`. Output shape:\n`[batch; seq_len; embed_dim]`.\n\nOptions:\n\n- `~num_kv_heads` — for grouped query attention (GQA). Default: same as\n `num_heads`.\n- `~is_causal:true` — applies a causal mask to prevent attending to\n future positions.\n- `~rope:true` — applies rotary position embeddings to Q and K.\n `~rope_theta` sets the base frequency (default 10000.0).\n- `~dropout` — attention dropout rate. Requires `~rngs` during training.\n\nPass an attention mask via `Context`:\n\n\n```ocaml\nlet ctx =\n Context.empty\n |> Context.set ~name:Attention.attention_mask_key (Ptree.P mask)\nin\nLayer.apply model vars ~training:false ~ctx input\n```\n\nThe mask is a bool or int32 tensor of shape `[batch; seq_k]`. Nonzero\npositions are kept, zero positions are masked.\n\nRoPE is also available as a standalone function:\n\n\n```ocaml\nlet x' = Attention.rope x (* default theta=10000, seq_dim=-2 *)\nlet x' = Attention.rope ~theta:500000. ~seq_dim:1 x\n```\n\n## Custom Layers\n\nA custom layer is a `{ init; apply }` record. Here is a residual block:\n\n\n```ocaml\nlet residual_block ~dim () : (float, float) Layer.t =\n let inner = Layer.sequential [\n Layer.linear ~in_features:dim ~out_features:dim ();\n Layer.relu ();\n Layer.linear ~in_features:dim ~out_features:dim ();\n ] in\n {\n init = inner.init;\n apply = (fun ~params ~state ~dtype ~training ?rngs ?ctx x ->\n let y, state' = inner.apply ~params ~state ~dtype ~training ?rngs ?ctx x in\n (Nx.add x y, state'));\n }\n```\n\nUse `Layer.make_vars` to build vars in custom `init` functions:\n\n\n```ocaml\nLayer.make_vars ~params ~state:Ptree.empty ~dtype\n```\n\n## Context\n\n`Context.t` carries per-call auxiliary data that specific layers read\nduring the forward pass. Most layers ignore it.\n\n\n```ocaml\nlet ctx =\n Context.empty\n |> Context.set ~name:\"attention_mask\" (Ptree.P mask)\n |> Context.set ~name:\"token_type_ids\" (Ptree.P ids)\nin\nLayer.apply model vars ~training:false ~ctx input_ids\n```\n\nContext is forwarded through `compose` and `sequential` to all sublayers.\n`Train.fit`, `Train.step`, and `Train.predict` accept an optional `~ctx`\nargument.\n\n## Weight Initialization\n\nOverride default initialization with `Init.t` values:\n\n\n```ocaml\nLayer.linear ~in_features:128 ~out_features:64\n ~weight_init:(Init.he_normal ())\n ~bias_init:Init.zeros\n ()\n```\n\nAvailable initializers:\n\n- `Init.zeros`, `Init.ones`, `Init.constant v`\n- `Init.uniform ~scale ()`, `Init.normal ~stddev ()`\n- `Init.glorot_uniform ()`, `Init.glorot_normal ()`\n- `Init.he_uniform ()`, `Init.he_normal ()`\n- `Init.lecun_uniform ()`, `Init.lecun_normal ()`\n- `Init.variance_scaling ~scale ~mode ~distribution ()`\n\n## Next Steps\n\n- [Training](../03-training/) — optimizers, losses, data pipelines, training loops\n- [Checkpoints and Pretrained Models](../04-checkpoints-and-pretrained/) — saving, loading, HuggingFace Hub\n"
},
{
"path": "packages/kaun/doc/03-training.md",
"content": "# Training\n\nThis guide covers optimizers, learning-rate schedules, loss functions,\ndata pipelines, the high-level training loop, metrics, and custom\ntraining.\n\n## Optimizers\n\nOptimizers are provided by the `Vega` package and take a `Schedule.t`\nas the learning rate:\n\n\n```ocaml\n(* SGD with momentum *)\nVega.sgd ~momentum:0.9 (Vega.Schedule.constant 0.1)\n\n(* Adam *)\nVega.adam (Vega.Schedule.constant 1e-3)\n\n(* AdamW with weight decay *)\nVega.adamw ~weight_decay:0.01 (Vega.Schedule.constant 1e-3)\n\n(* RMSprop *)\nVega.rmsprop (Vega.Schedule.constant 1e-3)\n\n(* Adagrad *)\nVega.adagrad (Vega.Schedule.constant 0.01)\n```\n\n`sgd` supports optional `~momentum` (default 0.0) and `~nesterov`\n(default false). `adam` and `adamw` support `~b1` (default 0.9), `~b2`\n(default 0.999), and `~eps` (default 1e-8). `rmsprop` supports `~decay`\n(default 0.9), `~eps`, and `~momentum`.\n\n## Learning-Rate Schedules\n\nA schedule is a function `int -> float` mapping 1-based step numbers to\nlearning rates:\n\n\n```ocaml\n(* Fixed learning rate *)\nVega.Schedule.constant 1e-3\n\n(* Cosine decay from 1e-3 to 0 over 10000 steps *)\nVega.Schedule.cosine_decay ~init_value:1e-3 ~decay_steps:10000 ()\n\n(* Cosine decay with minimum alpha *)\nVega.Schedule.cosine_decay ~init_value:1e-3 ~decay_steps:10000 ~alpha:1e-5 ()\n\n(* Linear warmup from 0 to 1e-3 over 1000 steps *)\nVega.Schedule.linear ~init_value:0. ~end_value:1e-3 ~steps:1000\n\n(* Cosine warmup *)\nVega.Schedule.warmup_cosine ~init_value:0. ~peak_value:1e-3 ~warmup_steps:1000\n\n(* Exponential decay *)\nVega.Schedule.exponential_decay ~init_value:1e-3 ~decay_rate:0.96 ~decay_steps:1000\n```\n\nCompose schedules by writing a custom function:\n\n\n```ocaml\nlet warmup_then_cosine step =\n if step <= 1000 then\n Vega.Schedule.linear ~init_value:0. ~end_value:1e-3 ~steps:1000 step\n else\n Vega.Schedule.cosine_decay ~init_value:1e-3 ~decay_steps:9000 () (step - 1000)\n```\n\n## Loss Functions\n\nAll loss functions return scalar tensors that are differentiable through\nRune's autodiff:\n\n\n```ocaml\n(* Multi-class: logits [batch; num_classes], one-hot labels [batch; num_classes] *)\nLoss.cross_entropy logits one_hot_labels\n\n(* Multi-class with integer labels: logits [batch; num_classes], labels [batch] *)\nLoss.cross_entropy_sparse logits class_indices\n\n(* Binary: raw logits (not sigmoid), labels in {0, 1} *)\nLoss.binary_cross_entropy logits labels\n\n(* Regression *)\nLoss.mse predictions targets\nLoss.mae predictions targets\n```\n\n## Data Pipelines\n\n`Data.t` is a lazy, composable iterator. Build pipelines by chaining\nconstructors, transformers, and consumers.\n\n### Constructors\n\n\n```ocaml\n(* From arrays *)\nData.of_array [| example1; example2; example3 |]\n\n(* From tensors: slices along first dimension *)\nData.of_tensor x (* yields x[0], x[1], ... *)\nData.of_tensors (x, y) (* yields (x[0], y[0]), (x[1], y[1]), ... *)\n\n(* From a function *)\nData.of_fn 1000 (fun i -> generate_example i)\n\n(* Repeat a value *)\nData.repeat 1000 (x, loss_fn)\n```\n\n### Transformers\n\n\n```ocaml\n(* Map each element *)\nData.map (fun (x, y) -> (preprocess x, y)) data\n\n(* Batch into arrays of size n *)\nData.batch 32 data (* yields arrays of 32 elements *)\nData.batch ~drop_last:true 32 data\n\n(* Batch and map in one step *)\nData.map_batch 32 collate_fn data\n\n(* Shuffle *)\nData.shuffle rng_key data\n```\n\n### Consumers\n\n\n```ocaml\nData.iter (fun x -> process x) data\nData.iteri (fun i x -> Printf.printf \"%d: %f\\n\" i x) data\nData.fold (fun acc x -> acc +. x) 0. data\nData.to_array data\nData.to_seq data\n```\n\n### The prepare Shortcut\n\n`Data.prepare` combines tensor slicing, optional shuffle, and batching\ninto one call. It is the standard way to feed tensor data to training:\n\n\n```ocaml\nlet train_data =\n Data.prepare ~shuffle:rng_key ~batch_size:64 (x_train, y_train)\n |> Data.map (fun (x, y) ->\n (x, fun logits -> Loss.cross_entropy_sparse logits y))\n```\n\n`Data.prepare` yields `(x_batch, y_batch)` tensor pairs. The `Data.map`\nstep attaches the loss function, producing the `(input, loss_fn)` pairs\nthat `Train.fit` expects.\n\n`~drop_last` defaults to `true` in `prepare`.\n\n### Resetting\n\nPipelines are single-pass. Call `Data.reset` to iterate again:\n\n\n```ocaml\nData.reset test_batches;\nlet acc = Metric.eval eval_fn test_batches\n```\n\n## High-Level Training\n\n### Train.make and Train.init\n\nCreate a trainer by pairing a model with an optimizer, then initialize:\n\n\n```ocaml\nlet trainer = Train.make ~model\n ~optimizer:(Vega.adam (Vega.Schedule.constant 1e-3))\n\nlet st = Train.init trainer ~dtype:Nx.Float32\n```\n\n### Train.fit\n\n`Train.fit` trains over a data pipeline and returns the final state:\n\n\n```ocaml\nlet st = Train.fit trainer st\n ~report:(fun ~step ~loss _st ->\n Printf.printf \"step %d loss %.4f\\n\" step loss)\n data\n```\n\nEach element of `data` is `(input, loss_fn)` where `loss_fn` takes the\nmodel output and returns a scalar loss.\n\nThe optional `~report` callback is called after every step. The `~step`\nnumber is 1-based.\n\n### Early Stopping\n\nRaise `Train.Early_stop` inside `~report` to end training early.\n`Train.fit` catches it and returns the current state:\n\n\n```ocaml\nlet st = Train.fit trainer st\n ~report:(fun ~step:_ ~loss st ->\n if loss < 0.001 then raise Train.Early_stop)\n data\n```\n\n### Train.predict\n\nRun inference in evaluation mode (no dropout, no state updates):\n\n\n```ocaml\nlet logits = Train.predict trainer st x\n```\n\n### Train.step\n\nFor manual control over single training steps:\n\n\n```ocaml\nlet loss, st' = Train.step trainer st ~training:true\n ~loss:(fun logits -> Loss.cross_entropy_sparse logits y)\n x\n```\n\n### Starting from Pretrained Weights\n\nUse `Train.make_state` to create training state from externally loaded\nweights instead of random initialization:\n\n\n```ocaml\nlet vars = (* load from checkpoint *) in\nlet st = Train.make_state trainer vars\n```\n\n## Metrics\n\n### Metric Functions\n\nMetric functions are plain `predictions -> targets -> float` functions:\n\n\n```ocaml\n(* Multi-class: logits [batch; num_classes], labels [batch] *)\nMetric.accuracy logits targets\n\n(* Binary classification *)\nMetric.binary_accuracy ~threshold:0.5 predictions targets\n\n(* Precision, recall, F1 with averaging mode *)\nMetric.precision Metric.Macro logits targets\nMetric.recall Metric.Micro logits targets\nMetric.f1 Metric.Weighted logits targets\n```\n\nAveraging modes: `Macro` (unweighted mean of per-class scores), `Micro`\n(global aggregation), `Weighted` (mean weighted by class support).\n\n### Dataset Evaluation\n\n`Metric.eval` folds a function over a data pipeline and returns the\nmean:\n\n\n```ocaml\nData.reset test_batches;\nlet test_acc =\n Metric.eval\n (fun (x, y) ->\n let logits = Train.predict trainer st x in\n Metric.accuracy logits y)\n test_batches\n```\n\n`Metric.eval_many` evaluates multiple named metrics at once:\n\n\n```ocaml\nlet results =\n Metric.eval_many\n (fun (x, y) ->\n let logits = Train.predict trainer st x in\n [ (\"accuracy\", Metric.accuracy logits y);\n (\"f1\", Metric.f1 Metric.Macro logits y) ])\n test_batches\n(* results : (string * float) list *)\n```\n\n### Running Tracker\n\n`Metric.tracker` accumulates running means during training:\n\n\n```ocaml\nlet tracker = Metric.tracker () in\n(* In the training loop: *)\nMetric.observe tracker \"loss\" loss_value;\nMetric.observe tracker \"accuracy\" acc_value;\n\n(* After an epoch: *)\nPrintf.printf \"%s\\n\" (Metric.summary tracker);\n(* \"accuracy: 0.9150 loss: 0.4231\" *)\n\nMetric.reset tracker\n```\n\n## Gradient Utilities\n\n### Gradient Clipping\n\nClip gradients by global L2 norm to prevent exploding gradients. Use\nthis with `Train.step` in custom training loops:\n\n\n```ocaml\nlet clipped_grads = Optim.clip_by_global_norm 1.0 grads\n```\n\n### Global Norm\n\nCompute the L2 norm across all leaf tensors:\n\n\n```ocaml\nlet norm = Optim.global_norm grads\n```\n\n### Manual Gradient Computation\n\n`Grad.value_and_grad` differentiates a function with respect to a\n`Ptree.t`:\n\n\n```ocaml\nlet loss, grads = Grad.value_and_grad\n (fun params ->\n let output, _state = model.apply ~params ~state ~dtype ~training:true x in\n Loss.mse output y)\n params\n```\n\n`Grad.value_and_grad_aux` returns auxiliary data alongside the loss:\n\n\n```ocaml\nlet loss, grads, new_state = Grad.value_and_grad_aux\n (fun params ->\n let output, new_state = model.apply ~params ~state ~dtype ~training:true x in\n (Loss.mse output y, new_state))\n params\n```\n\n## Next Steps\n\n- [Layers and Models](../02-layers-and-models/) — full layer catalog, composition, custom layers\n- [Checkpoints and Pretrained Models](../04-checkpoints-and-pretrained/) — saving, loading, HuggingFace Hub\n"
},
{
"path": "packages/kaun/doc/04-checkpoints-and-pretrained.md",
"content": "# Checkpoints and Pretrained Models\n\nThis guide covers saving and loading model parameters with SafeTensors,\nand downloading pretrained weights from the HuggingFace Hub.\n\n## SafeTensors Checkpointing\n\nKaun serializes parameter trees to the\n[SafeTensors](https://huggingface.co/docs/safetensors/) format. Tensor\npaths from the tree structure become file keys (e.g. `layers.0.weight`).\n\n### Saving\n\n\n```ocaml\nlet vars = Train.vars st in\nCheckpoint.save \"model.safetensors\" (Layer.params vars)\n```\n\n### Loading\n\n`Checkpoint.load` requires a `~like` template that defines the expected\ntree structure and dtypes. Tensors are cast to the template's dtype if\nneeded. Extra keys in the file are ignored.\n\n\n```ocaml\n(* Initialize model to get the tree structure *)\nlet vars = Layer.init model ~dtype:Nx.Float32 in\nlet params = Checkpoint.load \"model.safetensors\" ~like:(Layer.params vars) in\nlet vars = Layer.with_params vars params\n```\n\n### Saving and Loading State\n\nTo save both parameters and non-trainable state (e.g. batch norm\nrunning statistics):\n\n\n```ocaml\n(* Save *)\nlet vars = Train.vars st in\nCheckpoint.save \"params.safetensors\" (Layer.params vars);\nCheckpoint.save \"state.safetensors\" (Layer.state vars)\n\n(* Load *)\nlet vars = Layer.init model ~dtype:Nx.Float32 in\nlet params = Checkpoint.load \"params.safetensors\" ~like:(Layer.params vars) in\nlet state = Checkpoint.load \"state.safetensors\" ~like:(Layer.state vars) in\nlet vars = Layer.with_params vars params |> fun v -> Layer.with_state v state\n```\n\n### Resuming Training\n\nUse `Train.make_state` to create training state from loaded weights:\n\n\n```ocaml\nlet trainer = Train.make ~model ~optimizer in\nlet st = Train.make_state trainer vars in\n(* Continue training from here *)\nlet st = Train.fit trainer st data\n```\n\n## HuggingFace Hub\n\nThe `kaun-hf` package provides access to the HuggingFace Hub for\ndownloading pretrained model weights and configurations.\n\n### Downloading Files\n\n\n```ocaml\nlet path = Kaun_hf.download_file ~model_id:\"bert-base-uncased\"\n ~filename:\"config.json\" ()\n(* path : string — local filesystem path *)\n```\n\nFiles are cached under `$RAVEN_CACHE_ROOT/huggingface` (or\n`$XDG_CACHE_HOME/raven/huggingface`). Subsequent calls return the cached\npath.\n\nOptions:\n\n- `~token` — HuggingFace API token for private repositories. Defaults\n to the `HF_TOKEN` environment variable.\n- `~cache_dir` — override the default cache directory.\n- `~offline:true` — only return cached files, do not download.\n- `~revision:(Rev \"v1.0\")` — download a specific tag, branch, or commit.\n Default is `Main`.\n\n### Loading Configuration\n\n\n```ocaml\nlet config = Kaun_hf.load_config ~model_id:\"bert-base-uncased\" ()\n(* config : Jsont.json *)\n```\n\nReturns the parsed `config.json` from the repository.\n\n### Loading Weights\n\n\n```ocaml\nlet weights = Kaun_hf.load_weights ~model_id:\"bert-base-uncased\" ()\n(* weights : (string * Kaun.Ptree.tensor) list *)\n```\n\nReturns a flat list of `(name, tensor)` pairs from the model's\nSafeTensors checkpoint. Sharded checkpoints are handled transparently:\nwhen `model.safetensors.index.json` is present, all shards are\ndownloaded and merged.\n\nTensor names are the raw keys from the SafeTensors file (e.g.\n`bert.encoder.layer.0.attention.self.query.weight`). Your model code\nmaps these to its own parameter structure.\n\n### Loading a Pretrained Model\n\nThe typical pattern for loading pretrained weights:\n\n1. Build the model architecture from the config.\n2. Initialize to get the parameter tree structure.\n3. Load weights and map them to the tree.\n\n\n```ocaml\n(* 1. Build model from config *)\nlet config = Kaun_hf.load_config ~model_id:\"bert-base-uncased\" () in\nlet model = build_bert_model config in\n\n(* 2. Initialize to get tree structure *)\nlet vars = Layer.init model ~dtype:Nx.Float32 in\n\n(* 3. Load and map weights *)\nlet weights = Kaun_hf.load_weights ~model_id:\"bert-base-uncased\" () in\nlet params = map_weights_to_ptree weights (Layer.params vars) in\nlet vars = Layer.with_params vars params in\n\n(* 4. Use for inference *)\nlet trainer = Train.make ~model\n ~optimizer:(Vega.adam (Vega.Schedule.constant 1e-5))\nin\nlet st = Train.make_state trainer vars in\nlet logits = Train.predict trainer st input_ids\n```\n\n### Cache Management\n\n\n```ocaml\n(* Clear all cached files *)\nKaun_hf.clear_cache ()\n\n(* Clear a specific model's cache *)\nKaun_hf.clear_cache ~model_id:\"bert-base-uncased\" ()\n```\n\n## Next Steps\n\n- [Getting Started](../01-getting-started/) — XOR and MNIST examples\n- [Layers and Models](../02-layers-and-models/) — layer catalog, composition, custom layers\n- [Training](../03-training/) — optimizers, losses, data pipelines, training loops\n"
},
{
"path": "packages/kaun/doc/06-pytorch-comparison.md",
"content": "# Kaun vs. PyTorch / Flax -- A Practical Comparison\n\nThis guide explains how Kaun relates to PyTorch and Flax, focusing on:\n\n* How core concepts map (modules/layers, parameters, training loops)\n* Where the APIs feel similar vs. deliberately different\n* How to translate common patterns between frameworks\n\nKaun's design is closer to Flax than to PyTorch: layers are pure data,\nparameters are explicit trees, and forward passes are functions rather\nthan method calls. If you know Flax, Kaun will feel familiar. If you\nknow only PyTorch, the main shift is from mutable objects to immutable\nrecords.\n\n---\n\n## 1. Big-Picture Differences\n\n| Aspect | PyTorch | Flax (Linen) | Kaun (OCaml) |\n| ----------------- | ---------------------------------------------- | ---------------------------------- | ----------------------------------------------------- |\n| Language | Python, dynamic | Python (JAX), dynamic | OCaml, statically typed |\n| Model definition | `nn.Module` class with `forward` | `nn.Module` class with `__call__` | `Layer.t` record with `init` and `apply` |\n| Parameter storage | Mutable attributes on module | Frozen dict returned by `init` | `Ptree.t` tree returned by `Layer.init` |\n| Forward pass | `model(x)` (stateful method) | `model.apply(params, x)` | `Layer.apply model vars ~training x` |\n| Mutation | Modules are mutable objects | Params are immutable dicts | `Layer.vars` and `Ptree.t` are immutable |\n| Autograd | Dynamic tape (`autograd`) | Functional transforms (`jax.grad`) | Rune effect-based autodiff |\n| Optimizer | `torch.optim.Adam(model.parameters(), lr=...)` | `optax.adam(lr)` | `Vega.adam (Vega.Schedule.constant lr)` |\n| Training loop | Manual (or Lightning/etc.) | Manual (or Orbax/etc.) | `Train.fit` or manual `Train.step` |\n| Data loading | `DataLoader` | `tf.data` or manual | `Data.t` lazy pipeline |\n| Checkpointing | `torch.save` / `torch.load` (pickle) | Orbax / msgpack | SafeTensors via `Checkpoint.save` / `Checkpoint.load` |\n| RNG | Global `torch.manual_seed` | Explicit PRNGKey threading | Implicit scope via `Nx.Rng.run ~seed` |\n| Device management | `model.to(\"cuda\")`, `tensor.cuda()` | `jax.device_put` | CPU by default; JIT manages devices internally |\n\n---\n\n## 2. Defining Models\n\n### PyTorch\n\n```python\nimport torch\nimport torch.nn as nn\n\nclass MLP(nn.Module):\n def __init__(self, in_features, hidden, out_features):\n super().__init__()\n self.fc1 = nn.Linear(in_features, hidden)\n self.fc2 = nn.Linear(hidden, out_features)\n\n def forward(self, x):\n x = torch.relu(self.fc1(x))\n return self.fc2(x)\n\nmodel = MLP(784, 128, 10)\n```\n\n### Flax\n\n```python\nimport flax.linen as nn\nimport jax\n\nclass MLP(nn.Module):\n hidden: int\n out_features: int\n\n @nn.compact\n def __call__(self, x):\n x = nn.relu(nn.Dense(self.hidden)(x))\n return nn.Dense(self.out_features)(x)\n\nmodel = MLP(hidden=128, out_features=10)\nparams = model.init(jax.random.PRNGKey(0), jnp.ones([1, 784]))\n```\n\n### Kaun\n\n\n```ocaml\nopen Kaun\n\nlet model =\n Layer.sequential\n [\n Layer.linear ~in_features:784 ~out_features:128 ();\n Layer.relu ();\n Layer.linear ~in_features:128 ~out_features:10 ();\n ]\n\nlet vars =\n Nx.Rng.run ~seed:0 @@ fun () ->\n Layer.init model ~dtype:Nx.Float32\n```\n\nKey differences:\n\n* PyTorch defines models as classes. Flax defines models as dataclasses\n with `__call__`. Kaun uses `Layer.t` records -- plain data, not classes.\n* `Layer.sequential` replaces class-based composition for homogeneous\n float pipelines. `Layer.compose` handles heterogeneous types (e.g.\n embedding into dense).\n* Activation functions are layers (`Layer.relu ()`) rather than free\n functions called inside `forward`. This keeps the composition uniform.\n\n---\n\n## 3. Parameters\n\n### PyTorch\n\n```python\n# Parameters live inside the module\nfor name, param in model.named_parameters():\n print(name, param.shape)\n\n# state_dict is an OrderedDict\nsd = model.state_dict()\nmodel.load_state_dict(sd)\n```\n\n### Flax\n\n```python\n# Params are a frozen dict returned by init\nparams = model.init(key, x)[\"params\"]\njax.tree_util.tree_map(lambda p: p.shape, params)\n```\n\n### Kaun\n\n\n```ocaml\n(* vars bundles params, state, and dtype *)\nlet params = Layer.params vars (* Ptree.t *)\nlet state = Layer.state vars (* Ptree.t *)\nlet dt = Layer.dtype vars (* (float, 'layout) Nx.dtype *)\n\n(* Inspect parameter shapes *)\nlet paths = Ptree.flatten_with_paths params\n(* [(\"0.weight\", P tensor); (\"0.bias\", P tensor); ...] *)\n\n(* Count total parameters *)\nlet n = Ptree.count_parameters params\n\n(* Replace parameters *)\nlet vars' = Layer.with_params vars new_params\n```\n\nKey differences:\n\n* PyTorch stores parameters as mutable module attributes. Flax returns\n frozen dicts. Kaun returns `Ptree.t` -- a tree with `Tensor` leaves,\n `Dict` nodes, and `List` nodes.\n* `Ptree.t` is plain immutable data. You can map, fold, flatten, and\n serialize it without going through the model.\n* `Layer.vars` also carries non-trainable state (e.g. batch norm\n running statistics), separate from trainable parameters.\n\n---\n\n## 4. Forward Pass\n\n### PyTorch\n\n```python\nmodel.train()\noutput = model(x) # stateful: dropout active, batchnorm updates\n\nmodel.eval()\nwith torch.no_grad():\n output = model(x) # no dropout, batchnorm uses running stats\n```\n\n### Flax\n\n```python\noutput = model.apply(params, x)\noutput = model.apply(params, x, train=True, rngs={\"dropout\": key})\n```\n\n### Kaun\n\n\n```ocaml\n(* Training: dropout active, batchnorm updates running stats *)\nlet output, vars' = Layer.apply model vars ~training:true x\n\n(* Evaluation: no dropout, batchnorm uses running stats *)\nlet output, vars' = Layer.apply model vars ~training:false x\n\n(* Or through the trainer *)\nlet logits = Train.predict trainer st x\n```\n\nKey differences:\n\n* PyTorch uses `model.train()` / `model.eval()` to switch mode globally.\n Kaun passes `~training` as an argument on each call.\n* `Layer.apply` returns `(output, updated_vars)`. The updated vars carry\n new state (e.g. batch norm statistics). Parameters are unchanged.\n* `Train.predict` is a shortcut for evaluation mode with no state updates.\n\n---\n\n## 5. Optimizers and LR Schedules\n\n### PyTorch\n\n```python\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10000)\n\noptimizer.zero_grad()\nloss.backward()\noptimizer.step()\nscheduler.step()\n```\n\n### Flax / Optax\n\n```python\nimport optax\n\ntx = optax.adam(learning_rate=optax.cosine_decay_schedule(1e-3, 10000))\nopt_state = tx.init(params)\nupdates, opt_state = tx.update(grads, opt_state, params)\nparams = optax.apply_updates(params, updates)\n```\n\n### Kaun\n\n\n```ocaml\n(* Schedule is a function: step -> lr *)\nlet schedule = Vega.Schedule.cosine_decay ~init_value:1e-3 ~decay_steps:10000 ()\n\n(* Optimizer *)\nlet tx = Vega.adam schedule\n\n(* Init and update via Kaun's Optim bridge *)\nlet st = Optim.init tx params in\nlet updates, st' = Optim.update st params grads in\nlet params' = Optim.apply_updates params updates\n\n(* Or use the convenience function *)\nlet params', st' = Optim.step st params grads\n```\n\nAvailable optimizers:\n\n\n```ocaml\nVega.sgd ~momentum:0.9 ~nesterov:true schedule\nVega.adam ~b1:0.9 ~b2:0.999 ~eps:1e-8 schedule\nVega.adamw ~weight_decay:0.01 schedule\nVega.rmsprop ~decay:0.9 ~momentum:0.0 schedule\nVega.adagrad ~eps:1e-8 schedule\n```\n\nAvailable schedules:\n\n\n```ocaml\nVega.Schedule.constant 1e-3\nVega.Schedule.cosine_decay ~init_value:1e-3 ~decay_steps:10000 ()\nVega.Schedule.warmup_cosine ~init_value:0. ~peak_value:1e-3 ~warmup_steps:1000\nVega.Schedule.linear ~init_value:0. ~end_value:1e-3 ~steps:1000\nVega.Schedule.exponential_decay ~init_value:1e-3 ~decay_rate:0.96 ~decay_steps:1000\n```\n\nKey differences:\n\n* PyTorch couples the optimizer to the model via `model.parameters()`.\n Vega and Optax are decoupled -- they operate on parameter trees.\n* PyTorch separates scheduler from optimizer. Vega (like Optax) bakes\n the schedule into the optimizer as its last positional argument.\n* A Vega schedule is just `int -> float`. Compose them by writing a\n plain OCaml function.\n\n---\n\n## 6. Loss Functions\n\n### PyTorch\n\n```python\nloss = nn.functional.cross_entropy(logits, labels)\nloss = nn.functional.binary_cross_entropy_with_logits(logits, labels)\nloss = nn.functional.mse_loss(pred, target)\n```\n\n### Kaun\n\n\n```ocaml\n(* Multi-class with one-hot labels *)\nLoss.cross_entropy logits one_hot_labels\n\n(* Multi-class with integer labels *)\nLoss.cross_entropy_sparse logits class_indices\n\n(* Binary classification (raw logits, not sigmoid) *)\nLoss.binary_cross_entropy logits labels\n\n(* Regression *)\nLoss.mse predictions targets\nLoss.mae predictions targets\n```\n\nKey differences:\n\n* PyTorch's `cross_entropy` expects integer labels (like\n `cross_entropy_sparse`). Kaun offers both one-hot and integer\n variants.\n* All Kaun losses return scalar means and are differentiable through\n Rune's autodiff.\n* Kaun losses are plain functions, not module methods. There is no\n `nn.CrossEntropyLoss()` class.\n\n---\n\n## 7. Training Loops\n\n### PyTorch (manual loop)\n\n```python\nmodel.train()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\nfor epoch in range(10):\n for x_batch, y_batch in dataloader:\n optimizer.zero_grad()\n logits = model(x_batch)\n loss = nn.functional.cross_entropy(logits, y_batch)\n loss.backward()\n optimizer.step()\n print(f\"loss: {loss.item():.4f}\")\n```\n\n### Kaun with Train.fit\n\n\n```ocaml\nlet trainer =\n Train.make ~model\n ~optimizer:(Vega.adam (Vega.Schedule.constant 1e-3))\n\nlet st = Nx.Rng.run ~seed:42 @@ fun () ->\n Train.init trainer ~dtype:Nx.Float32\n\n(* Train over a data pipeline *)\nlet st =\n Train.fit trainer st\n ~report:(fun ~step ~loss _st ->\n Printf.printf \"step %d loss %.4f\\n\" step loss)\n data\n```\n\n`Train.fit` takes a `Data.t` where each element is `(input, loss_fn)`.\nThe loss function receives the model output and returns a scalar loss.\nGradient computation, optimizer step, and state threading are handled\ninternally.\n\n### Kaun with Train.step (manual loop)\n\nFor fine-grained control, use `Train.step` directly:\n\n\n```ocaml\nlet st = ref (Nx.Rng.run ~seed:42 @@ fun () ->\n Train.init trainer ~dtype:Nx.Float32)\n\nlet () =\n Data.iter\n (fun (x, y) ->\n let loss, st' =\n Train.step trainer !st ~training:true\n ~loss:(fun logits -> Loss.cross_entropy_sparse logits y)\n x\n in\n st := st';\n Printf.printf \"loss: %.4f\\n\" (Nx.item [] loss))\n data\n```\n\n### Early stopping\n\nRaise `Train.Early_stop` inside the `~report` callback:\n\n\n```ocaml\nlet st =\n Train.fit trainer st\n ~report:(fun ~step:_ ~loss _st ->\n if loss < 0.001 then raise Train.Early_stop)\n data\n```\n\nKey differences:\n\n* PyTorch training loops are fully manual: zero gradients, forward,\n backward, step. Kaun's `Train.fit` handles the entire loop.\n* `Train.step` is the escape hatch for custom loops, but you never call\n `backward` or `zero_grad` -- differentiation is implicit.\n* State threading replaces mutation. `Train.fit` returns the final\n state; `Train.step` returns `(loss, new_state)`.\n\n---\n\n## 8. Data Loading\n\n### PyTorch\n\n```python\nfrom torch.utils.data import DataLoader, TensorDataset\n\ndataset = TensorDataset(x_train, y_train)\nloader = DataLoader(dataset, batch_size=64, shuffle=True, drop_last=True)\n\nfor x_batch, y_batch in loader:\n ...\n```\n\n### Kaun\n\n\n```ocaml\n(* From tensor pairs -- the common case *)\nlet data =\n Data.prepare ~shuffle:true ~batch_size:64 (x_train, y_train)\n |> Data.map (fun (x, y) ->\n (x, fun logits -> Loss.cross_entropy_sparse logits y))\n\n(* From arrays *)\nlet data = Data.of_array examples |> Data.shuffle |> Data.batch 32\n\n(* From a generator function *)\nlet data = Data.of_fn 10000 generate_example\n\n(* Repeat a fixed example (useful for toy problems) *)\nlet data = Data.repeat 1000 (x, loss_fn)\n\n(* Consumers *)\nData.iter process data\nData.fold accumulate init data\nlet arr = Data.to_array data\n```\n\nKey differences:\n\n* PyTorch uses `Dataset` + `DataLoader` classes with worker processes\n for parallel loading. Kaun uses `Data.t`, a lazy composable iterator.\n* `Data.prepare` is the standard shortcut: it slices tensors, optionally\n shuffles, and batches in one call. `~drop_last` defaults to `true`.\n* Pipelines are single-pass. Call `Data.reset` before iterating again\n (e.g. between epochs).\n* `Data.map` attaches the loss function to each batch, producing the\n `(input, loss_fn)` pairs that `Train.fit` expects.\n\n---\n\n## 9. Checkpointing\n\n### PyTorch\n\n```python\n# Save\ntorch.save(model.state_dict(), \"model.pt\")\n\n# Load\nmodel.load_state_dict(torch.load(\"model.pt\"))\n```\n\n### Kaun\n\n\n```ocaml\n(* Save parameters *)\nlet vars = Train.vars st in\nCheckpoint.save \"model.safetensors\" (Layer.params vars)\n\n(* Load parameters *)\nlet vars = Layer.init model ~dtype:Nx.Float32 in\nlet params = Checkpoint.load \"model.safetensors\" ~like:(Layer.params vars) in\nlet vars = Layer.with_params vars params\n\n(* Save both params and state (e.g. batch norm stats) *)\nCheckpoint.save \"params.safetensors\" (Layer.params vars);\nCheckpoint.save \"state.safetensors\" (Layer.state vars)\n\n(* Resume training from loaded weights *)\nlet st = Train.make_state trainer vars\n```\n\nKey differences:\n\n* PyTorch uses Python pickle by default (arbitrary code execution risk).\n Kaun uses SafeTensors -- a flat, memory-mappable format with no code\n execution.\n* `Checkpoint.load` requires a `~like` template defining the expected\n tree structure and dtypes. Extra keys in the file are ignored, and\n tensors are cast to the template's dtype if needed.\n* Pretrained weights from HuggingFace Hub are available via\n `Kaun_hf.load_weights`.\n\n---\n\n## 10. Quick Cheat Sheet\n\n| Task | PyTorch | Kaun |\n| ---------------------------- | -------------------------------------------------- | -------------------------------------------------------------- |\n| Define a model | `class M(nn.Module): ...` | `Layer.sequential [Layer.linear ...; Layer.relu (); ...]` |\n| Initialize parameters | `model = M()` (implicit) | `Layer.init model ~dtype:Nx.Float32` |\n| Forward pass (training) | `model.train(); y = model(x)` | `Layer.apply model vars ~training:true x` |\n| Forward pass (eval) | `model.eval(); y = model(x)` | `Train.predict trainer st x` |\n| Count parameters | `sum(p.numel() for p in model.parameters())` | `Ptree.count_parameters (Layer.params vars)` |\n| Create optimizer | `Adam(model.parameters(), lr=1e-3)` | `Vega.adam (Vega.Schedule.constant 1e-3)` |\n| Cosine decay schedule | `CosineAnnealingLR(opt, T_max=N)` | `Vega.Schedule.cosine_decay ~init_value:lr ~decay_steps:N ()` |\n| Compute loss | `F.cross_entropy(logits, labels)` | `Loss.cross_entropy_sparse logits labels` |\n| Training step | `zero_grad(); loss.backward(); opt.step()` | `Train.step trainer st ~training:true ~loss x` |\n| Full training loop | Manual `for` loop | `Train.fit trainer st data` |\n| Early stopping | Manual condition check | `raise Train.Early_stop` inside `~report` |\n| Gradient clipping | `clip_grad_norm_(model.parameters(), max_norm)` | `Optim.clip_by_global_norm max_norm grads` |\n| Data loading | `DataLoader(dataset, batch_size=64, shuffle=True)` | `Data.prepare ~shuffle:true ~batch_size:64 (x, y)` |\n| Save checkpoint | `torch.save(model.state_dict(), path)` | `Checkpoint.save path (Layer.params vars)` |\n| Load checkpoint | `model.load_state_dict(torch.load(path))` | `Checkpoint.load path ~like:(Layer.params vars)` |\n| Compose heterogeneous layers | Define inside `forward` | `Layer.compose embedding_layer dense_layer` |\n| Dropout | `nn.Dropout(p=0.1)` | `Layer.dropout ~rate:0.1 ()` |\n| Batch normalization | `nn.BatchNorm2d(32)` | `Layer.batch_norm ~num_features:32 ()` |\n| Layer normalization | `nn.LayerNorm(128)` | `Layer.layer_norm ~dim:128 ()` |\n| Set RNG seed | `torch.manual_seed(42)` | `Nx.Rng.run ~seed:42 @@ fun () -> ...` |\n"
},
{
"path": "packages/kaun/doc/dune",
"content": "(mdx\n (files *.md)\n (package kaun)\n (libraries kaun rune nx nx.io))\n"
},
{
"path": "packages/kaun/doc/index.md",
"content": "# Kaun\n\nKaun is a neural network library for OCaml built on\n[Rune](https://github.com/raven-ml/raven/tree/main/rune). It provides\ncomposable layers, parameter trees, optimizers, data pipelines, and a\nhigh-level training loop. Pretrained weights load from the HuggingFace\nHub via SafeTensors.\n\n## Features\n\n- **Composable layers**: `sequential`, `compose`, and custom `{ init; apply }` records\n- **Parameter trees**: `Ptree.t` for inspection, serialization, and transformation\n- **High-level training**: `Train.fit` with data pipelines, or `Train.step` for manual control\n- **Optimizers**: SGD, Adam, AdamW, RMSprop, Adagrad with LR schedules\n- **Losses**: cross-entropy, binary cross-entropy, MSE, MAE\n- **Metrics**: accuracy, precision, recall, F1, running tracker, dataset evaluation\n- **Layers**: linear, conv1d/2d, layer_norm, rms_norm, batch_norm, embedding, dropout, pooling, multi-head attention with GQA and RoPE\n- **Checkpointing**: SafeTensors save/load, HuggingFace Hub integration\n- **Datasets**: MNIST and Fashion-MNIST loaders\n\n## Quick Start\n\nTrain a model on the XOR problem:\n\n\n```ocaml\nopen Kaun\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let x = Nx.create Nx.Float32 [| 4; 2 |] [| 0.; 0.; 0.; 1.; 1.; 0.; 1.; 1. |] in\n let y = Nx.create Nx.Float32 [| 4; 1 |] [| 0.; 1.; 1.; 0. |] in\n\n let model = Layer.sequential [\n Layer.linear ~in_features:2 ~out_features:4 ();\n Layer.tanh ();\n Layer.linear ~in_features:4 ~out_features:1 ();\n ] in\n\n let trainer = Train.make ~model\n ~optimizer:(Vega.adam (Vega.Schedule.constant 0.01))\n in\n let st = Train.init trainer ~dtype:Nx.Float32 in\n let st = Train.fit trainer st\n (Data.repeat 1000 (x, fun pred -> Loss.binary_cross_entropy pred y))\n in\n let pred = Train.predict trainer st x |> Nx.sigmoid in\n for i = 0 to 3 do\n Printf.printf \"[%.0f, %.0f] -> %.3f\\n\"\n (Nx.item [ i; 0 ] x) (Nx.item [ i; 1 ] x) (Nx.item [ i; 0 ] pred)\n done\n```\n\n## Next Steps\n\n- [Getting Started](01-getting-started/) — installation, XOR and MNIST examples, key concepts\n- [Layers and Models](02-layers-and-models/) — layer catalog, composition, custom layers\n- [Training](03-training/) — optimizers, losses, data pipelines, metrics, custom loops\n- [Checkpoints and Pretrained Models](04-checkpoints-and-pretrained/) — SafeTensors, HuggingFace Hub\n\n## See Also\n\n- [Munin](/docs/munin/) — experiment tracking with live terminal dashboard and CLI\n"
},
{
"path": "packages/kaun/examples/01-xor/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega kaun))\n"
},
{
"path": "packages/kaun/examples/01-xor/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Kaun\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let dtype = Nx.float32 in\n\n (* XOR dataset *)\n let x = Nx.create dtype [| 4; 2 |] [| 0.; 0.; 0.; 1.; 1.; 0.; 1.; 1. |] in\n let y = Nx.create dtype [| 4; 1 |] [| 0.; 1.; 1.; 0. |] in\n\n (* Model *)\n let model =\n Layer.sequential\n [\n Layer.linear ~in_features:2 ~out_features:4 ();\n Layer.tanh ();\n Layer.linear ~in_features:4 ~out_features:1 ();\n ]\n in\n\n (* Trainer = model + optimizer *)\n let trainer =\n Train.make ~model ~optimizer:(Vega.adam (Vega.Schedule.constant 0.01))\n in\n\n (* Initialize train state (vars + optimizer state) *)\n let st = Train.init trainer ~dtype in\n\n (* Fit *)\n let st =\n Train.fit trainer st\n ~report:(fun ~step ~loss _st ->\n if step mod 200 = 0 then Printf.printf \"step %4d loss %.6f\\n\" step loss)\n (Data.repeat 1000 (x, fun pred -> Loss.binary_cross_entropy pred y))\n in\n\n (* Evaluate *)\n let pred = Train.predict trainer st x |> Nx.sigmoid in\n Printf.printf \"\\npredictions (expected 0 1 1 0):\\n\";\n for i = 0 to 3 do\n Printf.printf \" [%.0f, %.0f] -> %.3f\\n\"\n (Nx.item [ i; 0 ] x)\n (Nx.item [ i; 1 ] x)\n (Nx.item [ i; 0 ] pred)\n done\n"
},
{
"path": "packages/kaun/examples/02-mnist/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega kaun kaun.datasets))\n"
},
{
"path": "packages/kaun/examples/02-mnist/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Kaun\n\nlet batch_size = 64\nlet epochs = 3\nlet lr = 0.001\n\nlet model =\n Layer.sequential\n [\n Layer.conv2d ~in_channels:1 ~out_channels:16 ();\n Layer.relu ();\n Layer.max_pool2d ~kernel_size:(2, 2) ();\n Layer.conv2d ~in_channels:16 ~out_channels:32 ();\n Layer.relu ();\n Layer.max_pool2d ~kernel_size:(2, 2) ();\n Layer.flatten ();\n Layer.linear ~in_features:(32 * 7 * 7) ~out_features:128 ();\n Layer.relu ();\n Layer.linear ~in_features:128 ~out_features:10 ();\n ]\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let dtype = Nx.float32 in\n\n Printf.printf \"Loading MNIST...\\n%!\";\n let (x_train, y_train), (x_test, y_test) = Kaun_datasets.mnist () in\n let n_train = (Nx.shape x_train).(0) in\n Printf.printf \" train: %d test: %d\\n%!\" n_train (Nx.shape x_test).(0);\n\n (* Test batches (fixed order, no shuffle) *)\n let test_batches = Data.prepare ~batch_size (x_test, y_test) in\n\n (* Trainer *)\n let trainer =\n Train.make ~model ~optimizer:(Vega.adam (Vega.Schedule.constant lr))\n in\n let st = ref (Train.init trainer ~dtype) in\n\n for epoch = 1 to epochs do\n let train_data =\n Data.prepare ~shuffle:true ~batch_size (x_train, y_train)\n |> Data.map (fun (x, y) ->\n (x, fun logits -> Loss.cross_entropy_sparse logits y))\n in\n let num_batches = n_train / batch_size in\n let tracker = Metric.tracker () in\n st :=\n Train.fit trainer !st\n ~report:(fun ~step ~loss _st ->\n Metric.observe tracker \"loss\" loss;\n Printf.printf \"\\r batch %d/%d loss: %.4f%!\" step num_batches loss)\n train_data;\n Printf.printf \"\\n%!\";\n\n (* Evaluate *)\n Data.reset test_batches;\n let test_acc =\n Metric.eval\n (fun (x, y) ->\n let logits = Train.predict trainer !st x in\n Metric.accuracy logits y)\n test_batches\n in\n\n Printf.printf \"epoch %d train_loss: %.4f test_acc: %.2f%%\\n%!\" epoch\n (Metric.mean tracker \"loss\")\n (test_acc *. 100.)\n done;\n\n Printf.printf \"\\nDone.\\n\"\n"
},
{
"path": "packages/kaun/examples/03-bert/bert.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Kaun\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\n\nlet require_float_dtype (type p in_elt) ~ctx (expected : (float, p) Nx.dtype)\n (x : (float, in_elt) Nx.t) : (float, p) Nx.t =\n match Nx_core.Dtype.equal_witness expected (Nx.dtype x) with\n | Some Type.Equal -> x\n | None ->\n invalid_argf \"%s: dtype mismatch (expected %s, got %s)\" ctx\n (Nx_core.Dtype.to_string expected)\n (Nx_core.Dtype.to_string (Nx.dtype x))\n\n(* Config *)\n\ntype config = {\n vocab_size : int;\n max_position_embeddings : int;\n type_vocab_size : int;\n hidden_size : int;\n num_hidden_layers : int;\n num_attention_heads : int;\n intermediate_size : int;\n hidden_dropout_prob : float;\n attention_dropout_prob : float;\n layer_norm_eps : float;\n}\n\nlet config ~vocab_size ~hidden_size ~num_hidden_layers ~num_attention_heads\n ~intermediate_size ?(max_position_embeddings = 512) ?(type_vocab_size = 2)\n ?(hidden_dropout_prob = 0.1) ?(attention_dropout_prob = 0.1)\n ?(layer_norm_eps = 1e-12) () =\n if hidden_size mod num_attention_heads <> 0 then\n invalid_argf\n \"Bert.config: hidden_size (%d) not divisible by num_attention_heads (%d)\"\n hidden_size num_attention_heads;\n if hidden_dropout_prob < 0.0 || hidden_dropout_prob >= 1.0 then\n invalid_arg \"Bert.config: hidden_dropout_prob must satisfy 0 <= p < 1\";\n if attention_dropout_prob < 0.0 || attention_dropout_prob >= 1.0 then\n invalid_arg \"Bert.config: attention_dropout_prob must satisfy 0 <= p < 1\";\n {\n vocab_size;\n max_position_embeddings;\n type_vocab_size;\n hidden_size;\n num_hidden_layers;\n num_attention_heads;\n intermediate_size;\n hidden_dropout_prob;\n attention_dropout_prob;\n layer_norm_eps;\n }\n\n(* Context keys *)\n\nlet token_type_ids_key = \"token_type_ids\"\n\n(* Helpers *)\n\nlet get_from_ctx_int32 ~name ~default ctx =\n match ctx with\n | Some c -> (\n match Context.find c ~name with\n | Some tensor -> Ptree.Tensor.to_typed_exn Nx.int32 tensor\n | None -> default ())\n | None -> default ()\n\nlet get_attention_mask_bool ctx ~batch ~seq =\n match ctx with\n | Some c -> (\n match Context.find c ~name:Attention.attention_mask_key with\n | Some tensor -> (\n match Ptree.Tensor.to_typed Nx.bool tensor with\n | Some m -> m\n | None ->\n let int_mask = Ptree.Tensor.to_typed_exn Nx.int32 tensor in\n Nx.not_equal int_mask (Nx.zeros Nx.int32 (Nx.shape int_mask)))\n | None -> Nx.broadcast_to [| batch; seq |] (Nx.scalar Nx.bool true))\n | None -> Nx.broadcast_to [| batch; seq |] (Nx.scalar Nx.bool true)\n\nlet fields ~ctx t = Ptree.Dict.fields_exn ~ctx t\nlet get fs ~name dtype = Ptree.Dict.get_tensor_exn fs ~name dtype\nlet find ~ctx key fs = Ptree.Dict.find_exn ~ctx key fs\n\n(* Self-attention with biased projections *)\n\nlet self_attention (type l) ~(cfg : config) ~(dtype : (float, l) Nx.dtype)\n ~training ~attention_mask ~params (x : (float, l) Nx.t) : (float, l) Nx.t =\n let shape = Nx.shape x in\n let batch = shape.(0) in\n let seq = shape.(1) in\n let h = cfg.hidden_size in\n let heads = cfg.num_attention_heads in\n let head_dim = h / heads in\n let fs = fields ~ctx:\"Bert.attention\" params in\n\n let proj name =\n let w = get fs ~name:(name ^ \"_weight\") dtype in\n let b = get fs ~name:(name ^ \"_bias\") dtype in\n fun t -> Nx.add (Nx.matmul t w) b\n in\n let q = proj \"q\" x in\n let k = proj \"k\" x in\n let v = proj \"v\" x in\n\n let split_heads t =\n Nx.reshape [| batch; seq; heads; head_dim |] t\n |> Nx.transpose ~axes:[ 0; 2; 1; 3 ]\n in\n let q = split_heads q in\n let k = split_heads k in\n let v = split_heads v in\n\n (* Broadcast mask [batch; seq] -> [batch; 1; 1; seq] *)\n let attention_mask = Nx.reshape [| batch; 1; 1; seq |] attention_mask in\n\n let dropout_rate =\n if training && cfg.attention_dropout_prob > 0.0 then\n Some cfg.attention_dropout_prob\n else None\n in\n\n let attn =\n Kaun.Fn.dot_product_attention ~attention_mask ?dropout_rate q k v\n in\n\n (* Merge heads *)\n let merged =\n Nx.transpose attn ~axes:[ 0; 2; 1; 3 ]\n |> Nx.contiguous\n |> Nx.reshape [| batch; seq; h |]\n in\n\n (* Output projection *)\n let o_w = get fs ~name:\"o_weight\" dtype in\n let o_b = get fs ~name:\"o_bias\" dtype in\n Nx.add (Nx.matmul merged o_w) o_b\n\n(* Encoder block *)\n\nlet encoder_block (type l) ~(cfg : config) ~(dtype : (float, l) Nx.dtype)\n ~training ~attention_mask ~params (x : (float, l) Nx.t) : (float, l) Nx.t =\n let fs = fields ~ctx:\"Bert.block\" params in\n\n (* Self-attention *)\n let attn_params = find ~ctx:\"Bert.block\" \"attention\" fs in\n let attn =\n self_attention ~cfg ~dtype ~training ~attention_mask ~params:attn_params x\n in\n\n (* Hidden dropout on attention output *)\n let attn =\n if training && cfg.hidden_dropout_prob > 0.0 then\n Kaun.Fn.dropout ~rate:cfg.hidden_dropout_prob attn\n else attn\n in\n\n (* Residual + LayerNorm (post-norm, original BERT) *)\n let ln1_g = get fs ~name:\"attn_ln_gamma\" dtype in\n let ln1_b = get fs ~name:\"attn_ln_beta\" dtype in\n let x =\n Kaun.Fn.layer_norm ~gamma:ln1_g ~beta:ln1_b ~epsilon:cfg.layer_norm_eps\n (Nx.add x attn)\n in\n\n (* FFN: up -> GELU -> down *)\n let ffn_up_w = get fs ~name:\"ffn_up_weight\" dtype in\n let ffn_up_b = get fs ~name:\"ffn_up_bias\" dtype in\n let ffn_down_w = get fs ~name:\"ffn_down_weight\" dtype in\n let ffn_down_b = get fs ~name:\"ffn_down_bias\" dtype in\n\n let y = Nx.add (Nx.matmul x ffn_up_w) ffn_up_b |> Kaun.Activation.gelu in\n let y = Nx.add (Nx.matmul y ffn_down_w) ffn_down_b in\n\n (* Hidden dropout on FFN output *)\n let y =\n if training && cfg.hidden_dropout_prob > 0.0 then\n Kaun.Fn.dropout ~rate:cfg.hidden_dropout_prob y\n else y\n in\n\n (* Residual + LayerNorm *)\n let ln2_g = get fs ~name:\"ffn_ln_gamma\" dtype in\n let ln2_b = get fs ~name:\"ffn_ln_beta\" dtype in\n Kaun.Fn.layer_norm ~gamma:ln2_g ~beta:ln2_b ~epsilon:cfg.layer_norm_eps\n (Nx.add x y)\n\n(* Forward: embeddings + encoder stack *)\n\nlet encode (type l in_elt) ~(cfg : config) ~params\n ~(dtype : (float, l) Nx.dtype) ~training ?ctx\n (input_ids : (int32, in_elt) Nx.t) : (float, l) Nx.t =\n let input_ids = Nx.cast Nx.int32 input_ids in\n let shape = Nx.shape input_ids in\n let batch = shape.(0) in\n let seq = shape.(1) in\n\n if seq > cfg.max_position_embeddings then\n invalid_argf \"Bert.encode: seq_len=%d exceeds max_position_embeddings=%d\"\n seq cfg.max_position_embeddings;\n\n (* Read auxiliary inputs from context *)\n let token_type_ids =\n get_from_ctx_int32 ~name:token_type_ids_key ctx ~default:(fun () ->\n Nx.zeros Nx.int32 [| batch; seq |])\n in\n let attention_mask = get_attention_mask_bool ctx ~batch ~seq in\n\n (* Params *)\n let root = fields ~ctx:\"Bert.encode\" params in\n let emb_t = find ~ctx:\"Bert.encode\" \"embeddings\" root in\n let layers_t = find ~ctx:\"Bert.encode\" \"layers\" root in\n\n let emb = fields ~ctx:\"Bert.embeddings\" emb_t in\n let word_emb = get emb ~name:\"word\" dtype in\n let pos_emb = get emb ~name:\"pos\" dtype in\n let type_emb = get emb ~name:\"type\" dtype in\n let ln_g = get emb ~name:\"ln_gamma\" dtype in\n let ln_b = get emb ~name:\"ln_beta\" dtype in\n\n (* Embedding lookup: word + position + token_type *)\n let position_ids =\n Nx.arange_f Nx.float32 0.0 (float_of_int seq) 1.0\n |> Nx.cast Nx.int32\n |> Nx.reshape [| 1; seq |]\n |> Nx.broadcast_to [| batch; seq |]\n |> Nx.contiguous\n in\n let token_type_ids = Nx.contiguous token_type_ids in\n let tok = Kaun.Fn.embedding ~scale:false ~embedding:word_emb input_ids in\n let pos = Kaun.Fn.embedding ~scale:false ~embedding:pos_emb position_ids in\n let typ = Kaun.Fn.embedding ~scale:false ~embedding:type_emb token_type_ids in\n let x = Nx.add tok (Nx.add pos typ) in\n let x =\n Kaun.Fn.layer_norm ~gamma:ln_g ~beta:ln_b ~epsilon:cfg.layer_norm_eps x\n in\n\n (* Embedding dropout *)\n let x =\n if training && cfg.hidden_dropout_prob > 0.0 then\n Kaun.Fn.dropout ~rate:cfg.hidden_dropout_prob x\n else x\n in\n\n (* Encoder stack *)\n let blocks = Ptree.List.items_exn ~ctx:\"Bert.encode.layers\" layers_t in\n let x =\n List.fold_left\n (fun h block_params ->\n encoder_block ~cfg ~dtype ~training ~attention_mask ~params:block_params\n h)\n x blocks\n in\n x\n\n(* Parameter initialization *)\n\nlet init_block_params ~dtype ~hidden ~intermediate =\n let w = Init.normal ~stddev:0.02 () in\n let zeros n = Nx.zeros dtype [| n |] in\n let ones n = Nx.ones dtype [| n |] in\n let attn_params =\n Ptree.dict\n [\n (\"q_weight\", Ptree.tensor (w.f [| hidden; hidden |] dtype));\n (\"q_bias\", Ptree.tensor (zeros hidden));\n (\"k_weight\", Ptree.tensor (w.f [| hidden; hidden |] dtype));\n (\"k_bias\", Ptree.tensor (zeros hidden));\n (\"v_weight\", Ptree.tensor (w.f [| hidden; hidden |] dtype));\n (\"v_bias\", Ptree.tensor (zeros hidden));\n (\"o_weight\", Ptree.tensor (w.f [| hidden; hidden |] dtype));\n (\"o_bias\", Ptree.tensor (zeros hidden));\n ]\n in\n Ptree.dict\n [\n (\"attention\", attn_params);\n (\"attn_ln_gamma\", Ptree.tensor (ones hidden));\n (\"attn_ln_beta\", Ptree.tensor (zeros hidden));\n (\"ffn_up_weight\", Ptree.tensor (w.f [| hidden; intermediate |] dtype));\n (\"ffn_up_bias\", Ptree.tensor (zeros intermediate));\n (\"ffn_down_weight\", Ptree.tensor (w.f [| intermediate; hidden |] dtype));\n (\"ffn_down_bias\", Ptree.tensor (zeros hidden));\n (\"ffn_ln_gamma\", Ptree.tensor (ones hidden));\n (\"ffn_ln_beta\", Ptree.tensor (zeros hidden));\n ]\n\nlet init_encoder_params ~cfg ~dtype =\n let h = cfg.hidden_size in\n let w = Init.normal ~stddev:0.02 () in\n let word = w.f [| cfg.vocab_size; h |] dtype in\n let pos = w.f [| cfg.max_position_embeddings; h |] dtype in\n let typ = w.f [| cfg.type_vocab_size; h |] dtype in\n let blocks =\n List.init cfg.num_hidden_layers (fun _ ->\n init_block_params ~dtype ~hidden:h ~intermediate:cfg.intermediate_size)\n in\n Ptree.dict\n [\n ( \"embeddings\",\n Ptree.dict\n [\n (\"word\", Ptree.tensor word);\n (\"pos\", Ptree.tensor pos);\n (\"type\", Ptree.tensor typ);\n (\"ln_gamma\", Ptree.tensor (Nx.ones dtype [| h |]));\n (\"ln_beta\", Ptree.tensor (Nx.zeros dtype [| h |]));\n ] );\n (\"layers\", Ptree.list blocks);\n ]\n\n(* Layers *)\n\nlet encoder (cfg : config) () : (int32, float) Layer.t =\n {\n Layer.init =\n (fun ~dtype ->\n Layer.make_vars\n ~params:(init_encoder_params ~cfg ~dtype)\n ~state:Ptree.empty ~dtype);\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore state;\n let y = encode ~cfg ~params ~dtype ~training ?ctx x in\n (y, Ptree.empty));\n }\n\nlet pooler (cfg : config) () : (float, float) Layer.t =\n let w_init = Init.normal ~stddev:0.02 () in\n {\n Layer.init =\n (fun ~dtype ->\n let w = w_init.f [| cfg.hidden_size; cfg.hidden_size |] dtype in\n let b = Nx.zeros dtype [| cfg.hidden_size |] in\n Layer.make_vars\n ~params:\n (Ptree.dict\n [ (\"weight\", Ptree.tensor w); (\"bias\", Ptree.tensor b) ])\n ~state:Ptree.empty ~dtype);\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (training, ctx, state);\n let x = require_float_dtype ~ctx:\"Bert.pooler\" dtype x in\n let fs = fields ~ctx:\"Bert.pooler\" params in\n let w = get fs ~name:\"weight\" dtype in\n let b = get fs ~name:\"bias\" dtype in\n let batch = (Nx.shape x).(0) in\n let cls =\n Nx.slice [ A; R (0, 1) ] x |> Nx.reshape [| batch; cfg.hidden_size |]\n in\n (Nx.add (Nx.matmul cls w) b |> Nx.tanh, Ptree.empty));\n }\n\nlet for_sequence_classification (cfg : config) ~num_labels () :\n (int32, float) Layer.t =\n let w_init = Init.normal ~stddev:0.02 () in\n {\n Layer.init =\n (fun ~dtype ->\n let enc = init_encoder_params ~cfg ~dtype in\n let pool_w = w_init.f [| cfg.hidden_size; cfg.hidden_size |] dtype in\n let cls_w = w_init.f [| cfg.hidden_size; num_labels |] dtype in\n Layer.make_vars\n ~params:\n (Ptree.dict\n [\n (\"encoder\", enc);\n ( \"pooler\",\n Ptree.dict\n [\n (\"weight\", Ptree.tensor pool_w);\n ( \"bias\",\n Ptree.tensor (Nx.zeros dtype [| cfg.hidden_size |]) );\n ] );\n ( \"classifier\",\n Ptree.dict\n [\n (\"weight\", Ptree.tensor cls_w);\n (\"bias\", Ptree.tensor (Nx.zeros dtype [| num_labels |]));\n ] );\n ])\n ~state:Ptree.empty ~dtype);\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore state;\n let root = fields ~ctx:\"Bert.seq_cls\" params in\n let enc_params = find ~ctx:\"Bert.seq_cls\" \"encoder\" root in\n let pool_params = find ~ctx:\"Bert.seq_cls\" \"pooler\" root in\n let cls_params = find ~ctx:\"Bert.seq_cls\" \"classifier\" root in\n\n let hidden = encode ~cfg ~params:enc_params ~dtype ~training ?ctx x in\n\n (* Pooler: CLS token -> dense -> tanh *)\n let pool_fs = fields ~ctx:\"Bert.seq_cls.pooler\" pool_params in\n let pool_w = get pool_fs ~name:\"weight\" dtype in\n let pool_b = get pool_fs ~name:\"bias\" dtype in\n let batch = (Nx.shape hidden).(0) in\n let cls =\n Nx.slice [ A; R (0, 1) ] hidden\n |> Nx.reshape [| batch; cfg.hidden_size |]\n in\n let pooled = Nx.add (Nx.matmul cls pool_w) pool_b |> Nx.tanh in\n\n (* Dropout on pooled output during fine-tuning *)\n let pooled =\n if training && cfg.hidden_dropout_prob > 0.0 then\n Kaun.Fn.dropout ~rate:cfg.hidden_dropout_prob pooled\n else pooled\n in\n\n (* Classifier *)\n let cls_fs = fields ~ctx:\"Bert.seq_cls.classifier\" cls_params in\n let cls_w = get cls_fs ~name:\"weight\" dtype in\n let cls_b = get cls_fs ~name:\"bias\" dtype in\n (Nx.add (Nx.matmul pooled cls_w) cls_b, Ptree.empty));\n }\n\nlet for_masked_lm (cfg : config) () : (int32, float) Layer.t =\n let w_init = Init.normal ~stddev:0.02 () in\n {\n Layer.init =\n (fun ~dtype ->\n let enc = init_encoder_params ~cfg ~dtype in\n let dense_w = w_init.f [| cfg.hidden_size; cfg.hidden_size |] dtype in\n Layer.make_vars\n ~params:\n (Ptree.dict\n [\n (\"encoder\", enc);\n ( \"mlm\",\n Ptree.dict\n [\n (\"dense_weight\", Ptree.tensor dense_w);\n ( \"dense_bias\",\n Ptree.tensor (Nx.zeros dtype [| cfg.hidden_size |]) );\n ( \"ln_gamma\",\n Ptree.tensor (Nx.ones dtype [| cfg.hidden_size |]) );\n ( \"ln_beta\",\n Ptree.tensor (Nx.zeros dtype [| cfg.hidden_size |]) );\n ( \"decoder_bias\",\n Ptree.tensor (Nx.zeros dtype [| cfg.vocab_size |]) );\n ] );\n ])\n ~state:Ptree.empty ~dtype);\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore state;\n let root = fields ~ctx:\"Bert.mlm\" params in\n let enc_params = find ~ctx:\"Bert.mlm\" \"encoder\" root in\n let mlm_params = find ~ctx:\"Bert.mlm\" \"mlm\" root in\n\n let hidden = encode ~cfg ~params:enc_params ~dtype ~training ?ctx x in\n\n (* MLM transform: dense -> GELU -> LN *)\n let mlm_fs = fields ~ctx:\"Bert.mlm.head\" mlm_params in\n let dw = get mlm_fs ~name:\"dense_weight\" dtype in\n let db = get mlm_fs ~name:\"dense_bias\" dtype in\n let ln_g = get mlm_fs ~name:\"ln_gamma\" dtype in\n let ln_b = get mlm_fs ~name:\"ln_beta\" dtype in\n let dec_b = get mlm_fs ~name:\"decoder_bias\" dtype in\n\n let h = Nx.add (Nx.matmul hidden dw) db |> Kaun.Activation.gelu in\n let h =\n Kaun.Fn.layer_norm ~gamma:ln_g ~beta:ln_b ~epsilon:cfg.layer_norm_eps\n h\n in\n\n (* Tied decoder: logits = h @ word_emb^T + bias *)\n let enc_root = fields ~ctx:\"Bert.mlm.encoder\" enc_params in\n let emb_t = find ~ctx:\"Bert.mlm.encoder\" \"embeddings\" enc_root in\n let emb_fs = fields ~ctx:\"Bert.mlm.embeddings\" emb_t in\n let word_emb = get emb_fs ~name:\"word\" dtype in\n let logits =\n Nx.add (Nx.matmul h (Nx.transpose word_emb ~axes:[ 1; 0 ])) dec_b\n in\n (logits, Ptree.empty));\n }\n\n(* JSON config parsing *)\n\nlet json_mem name = function\n | Jsont.Object (mems, _) -> (\n match Jsont.Json.find_mem name mems with\n | Some (_, v) -> v\n | None -> Jsont.Null ((), Jsont.Meta.none))\n | _ -> Jsont.Null ((), Jsont.Meta.none)\n\nlet json_to_int = function\n | Jsont.Number (f, _) -> int_of_float f\n | _ -> failwith \"expected int\"\n\nlet json_to_int_option = function\n | Jsont.Number (f, _) -> Some (int_of_float f)\n | _ -> None\n\nlet json_to_float_option = function Jsont.Number (f, _) -> Some f | _ -> None\n\nlet parse_config json =\n config\n ~vocab_size:(json |> json_mem \"vocab_size\" |> json_to_int)\n ~hidden_size:(json |> json_mem \"hidden_size\" |> json_to_int)\n ~num_hidden_layers:(json |> json_mem \"num_hidden_layers\" |> json_to_int)\n ~num_attention_heads:(json |> json_mem \"num_attention_heads\" |> json_to_int)\n ~intermediate_size:(json |> json_mem \"intermediate_size\" |> json_to_int)\n ?max_position_embeddings:\n (json |> json_mem \"max_position_embeddings\" |> json_to_int_option)\n ?type_vocab_size:(json |> json_mem \"type_vocab_size\" |> json_to_int_option)\n ?hidden_dropout_prob:\n (json |> json_mem \"hidden_dropout_prob\" |> json_to_float_option)\n ?attention_dropout_prob:\n (json |> json_mem \"attention_probs_dropout_prob\" |> json_to_float_option)\n ?layer_norm_eps:(json |> json_mem \"layer_norm_eps\" |> json_to_float_option)\n ()\n\n(* HuggingFace weight mapping *)\n\nlet transpose_weight (Ptree.P t) = Ptree.P (Nx.transpose t ~axes:[ 1; 0 ])\nlet cast_tensor dtype (Ptree.P t) = Ptree.P (Nx.cast dtype t)\n\nlet map_hf_weights ~cfg ~dtype hf_weights =\n let tbl = Hashtbl.create (List.length hf_weights) in\n List.iter (fun (name, tensor) -> Hashtbl.add tbl name tensor) hf_weights;\n let hf name =\n match Hashtbl.find_opt tbl name with\n | Some t -> cast_tensor dtype t\n | None -> invalid_argf \"from_pretrained: missing HF weight %S\" name\n in\n (* Some checkpoints use LayerNorm.weight/bias, others use\n LayerNorm.gamma/beta. Try both. *)\n let hf_ln_weight prefix =\n let w = prefix ^ \".weight\" in\n let g = prefix ^ \".gamma\" in\n if Hashtbl.mem tbl w then hf w else hf g\n in\n let hf_ln_bias prefix =\n let b = prefix ^ \".bias\" in\n let beta = prefix ^ \".beta\" in\n if Hashtbl.mem tbl b then hf b else hf beta\n in\n let hf_t name = Ptree.Tensor (transpose_weight (hf name)) in\n let hf_b name = Ptree.Tensor (hf name) in\n let ln_w prefix = Ptree.Tensor (hf_ln_weight prefix) in\n let ln_b prefix = Ptree.Tensor (hf_ln_bias prefix) in\n let layer i =\n let p s = Printf.sprintf \"bert.encoder.layer.%d.%s\" i s in\n let attn_ln = p \"attention.output.LayerNorm\" in\n let ffn_ln = p \"output.LayerNorm\" in\n Ptree.dict\n [\n ( \"attention\",\n Ptree.dict\n [\n (\"q_weight\", hf_t (p \"attention.self.query.weight\"));\n (\"q_bias\", hf_b (p \"attention.self.query.bias\"));\n (\"k_weight\", hf_t (p \"attention.self.key.weight\"));\n (\"k_bias\", hf_b (p \"attention.self.key.bias\"));\n (\"v_weight\", hf_t (p \"attention.self.value.weight\"));\n (\"v_bias\", hf_b (p \"attention.self.value.bias\"));\n (\"o_weight\", hf_t (p \"attention.output.dense.weight\"));\n (\"o_bias\", hf_b (p \"attention.output.dense.bias\"));\n ] );\n (\"attn_ln_gamma\", ln_w attn_ln);\n (\"attn_ln_beta\", ln_b attn_ln);\n (\"ffn_up_weight\", hf_t (p \"intermediate.dense.weight\"));\n (\"ffn_up_bias\", hf_b (p \"intermediate.dense.bias\"));\n (\"ffn_down_weight\", hf_t (p \"output.dense.weight\"));\n (\"ffn_down_bias\", hf_b (p \"output.dense.bias\"));\n (\"ffn_ln_gamma\", ln_w ffn_ln);\n (\"ffn_ln_beta\", ln_b ffn_ln);\n ]\n in\n let emb_ln = \"bert.embeddings.LayerNorm\" in\n let encoder_params =\n Ptree.dict\n [\n ( \"embeddings\",\n Ptree.dict\n [\n (\"word\", hf_b \"bert.embeddings.word_embeddings.weight\");\n (\"pos\", hf_b \"bert.embeddings.position_embeddings.weight\");\n (\"type\", hf_b \"bert.embeddings.token_type_embeddings.weight\");\n (\"ln_gamma\", ln_w emb_ln);\n (\"ln_beta\", ln_b emb_ln);\n ] );\n (\"layers\", Ptree.list (List.init cfg.num_hidden_layers layer));\n ]\n in\n let pooler_params =\n let has_pooler = Hashtbl.mem tbl \"bert.pooler.dense.weight\" in\n if has_pooler then\n Some\n (Ptree.dict\n [\n (\"weight\", hf_t \"bert.pooler.dense.weight\");\n (\"bias\", hf_b \"bert.pooler.dense.bias\");\n ])\n else None\n in\n let mlm_params =\n let has_mlm = Hashtbl.mem tbl \"cls.predictions.transform.dense.weight\" in\n if has_mlm then\n let mlm_ln = \"cls.predictions.transform.LayerNorm\" in\n Some\n (Ptree.dict\n [\n (\"dense_weight\", hf_t \"cls.predictions.transform.dense.weight\");\n (\"dense_bias\", hf_b \"cls.predictions.transform.dense.bias\");\n (\"ln_gamma\", ln_w mlm_ln);\n (\"ln_beta\", ln_b mlm_ln);\n (\"decoder_bias\", hf_b \"cls.predictions.bias\");\n ])\n else None\n in\n (encoder_params, pooler_params, mlm_params)\n\n(* Pretrained loading *)\n\nlet from_pretrained ?(model_id = \"bert-base-uncased\") () =\n let json = Kaun_hf.load_config ~model_id () in\n let cfg = parse_config json in\n let hf_weights = Kaun_hf.load_weights ~model_id () in\n let encoder_params, pooler_params, mlm_params =\n map_hf_weights ~cfg ~dtype:Nx.float32 hf_weights\n in\n (cfg, encoder_params, pooler_params, mlm_params)\n"
},
{
"path": "packages/kaun/examples/03-bert/bert.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** BERT encoder and task heads.\n\n BERT inputs are passed as int32 [input_ids] with auxiliary data in\n {!Context}:\n\n {[\n let ctx =\n Context.empty\n |> Context.set ~name:Bert.token_type_ids_key (Ptree.P token_type_ids)\n |> Context.set ~name:Attention.attention_mask_key (Ptree.P attention_mask)\n in\n Layer.apply model vars ~training:false ~ctx input_ids\n ]}\n\n When absent, [token_type_ids] defaults to zeros and [attention_mask]\n defaults to ones (no padding). *)\n\nopen Kaun\n\n(** {1:config Configuration} *)\n\ntype config = {\n vocab_size : int;\n max_position_embeddings : int;\n type_vocab_size : int;\n hidden_size : int;\n num_hidden_layers : int;\n num_attention_heads : int;\n intermediate_size : int;\n hidden_dropout_prob : float;\n attention_dropout_prob : float;\n layer_norm_eps : float;\n}\n(** The type for BERT configurations. *)\n\nval config :\n vocab_size:int ->\n hidden_size:int ->\n num_hidden_layers:int ->\n num_attention_heads:int ->\n intermediate_size:int ->\n ?max_position_embeddings:int ->\n ?type_vocab_size:int ->\n ?hidden_dropout_prob:float ->\n ?attention_dropout_prob:float ->\n ?layer_norm_eps:float ->\n unit ->\n config\n(** [config ~vocab_size ~hidden_size ~num_hidden_layers ~num_attention_heads\n ~intermediate_size ()] is a BERT configuration.\n\n [max_position_embeddings] defaults to [512]. [type_vocab_size] defaults to\n [2]. [hidden_dropout_prob] and [attention_dropout_prob] default to [0.1].\n [layer_norm_eps] defaults to [1e-12].\n\n Raises [Invalid_argument] if [hidden_size] is not divisible by\n [num_attention_heads] or if dropout rates are outside [\\[0, 1)]. *)\n\n(** {1:context Context keys} *)\n\nval token_type_ids_key : string\n(** [\"token_type_ids\"]. The {!Context} key for segment ids (shape\n [[batch; seq]], int32, values 0 or 1). *)\n\n(** {1:layers Layers} *)\n\nval encoder : config -> unit -> (int32, float) Layer.t\n(** [encoder cfg ()] is the base BERT encoder.\n\n Input: int32 [input_ids] of shape [[batch; seq]]. Output: float hidden\n states of shape [[batch; seq; hidden_size]].\n\n Reads {!token_type_ids_key} and {!Attention.attention_mask_key} from [ctx].\n*)\n\nval pooler : config -> unit -> (float, float) Layer.t\n(** [pooler cfg ()] maps [[batch; seq; hidden_size]] to [[batch; hidden_size]]\n by extracting the CLS token (position 0) and applying a dense + tanh. *)\n\nval for_sequence_classification :\n config -> num_labels:int -> unit -> (int32, float) Layer.t\n(** [for_sequence_classification cfg ~num_labels ()] is encoder + pooler +\n classifier. Output: logits [[batch; num_labels]]. *)\n\nval for_masked_lm : config -> unit -> (int32, float) Layer.t\n(** [for_masked_lm cfg ()] is encoder + MLM head with tied word embeddings.\n Output: logits [[batch; seq; vocab_size]]. *)\n\n(** {1:pretrained Pretrained loading} *)\n\nval from_pretrained :\n ?model_id:string -> unit -> config * Ptree.t * Ptree.t option * Ptree.t option\n(** [from_pretrained ?model_id ()] downloads [model_id] from HuggingFace and\n returns [(cfg, encoder_params, pooler_params, mlm_params)].\n\n [encoder_params] is ready for {!encoder}. [pooler_params] and [mlm_params]\n are [Some _] when the checkpoint contains the corresponding weights.\n\n [model_id] defaults to [\"bert-base-uncased\"]. *)\n"
},
{
"path": "packages/kaun/examples/03-bert/dune",
"content": "(executable\n (name main)\n (libraries nx nx.core rune vega kaun kaun.hf jsont))\n"
},
{
"path": "packages/kaun/examples/03-bert/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Fine-tune pretrained BERT for binary sentiment classification.\n\n Downloads bert-base-uncased from HuggingFace (~440MB on first run), assembles\n a sequence-classification head, and trains on a tiny synthetic dataset to\n show the full pipeline. *)\n\nopen Kaun\n\nlet print_shape name t =\n let shape = Nx.shape t in\n Printf.printf \"%s: [%s]\\n\" name\n (String.concat \"; \" (Array.to_list (Array.map string_of_int shape)))\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let dtype = Nx.float32 in\n let num_labels = 2 in\n\n (* Load pretrained encoder + pooler from HuggingFace *)\n Printf.printf \"Loading bert-base-uncased...\\n%!\";\n let cfg, encoder_params, pooler_params, _mlm_params =\n Bert.from_pretrained ()\n in\n Printf.printf \" hidden=%d layers=%d heads=%d vocab=%d\\n\\n\" cfg.hidden_size\n cfg.num_hidden_layers cfg.num_attention_heads cfg.vocab_size;\n\n (* Assemble classification model: pretrained encoder + pooler, fresh\n classifier head *)\n let w_init = Init.normal ~stddev:0.02 () in\n let params =\n Ptree.dict\n [\n (\"encoder\", encoder_params);\n ( \"pooler\",\n match pooler_params with\n | Some p -> p\n | None ->\n Ptree.dict\n [\n ( \"weight\",\n Ptree.tensor\n (w_init.f [| cfg.hidden_size; cfg.hidden_size |] dtype) );\n (\"bias\", Ptree.tensor (Nx.zeros dtype [| cfg.hidden_size |]));\n ] );\n ( \"classifier\",\n Ptree.dict\n [\n ( \"weight\",\n Ptree.tensor (w_init.f [| cfg.hidden_size; num_labels |] dtype)\n );\n (\"bias\", Ptree.tensor (Nx.zeros dtype [| num_labels |]));\n ] );\n ]\n in\n\n let model = Bert.for_sequence_classification cfg ~num_labels () in\n let vars = Layer.make_vars ~params ~state:Ptree.empty ~dtype in\n\n (* Tiny synthetic dataset (token ids from bert-base-uncased tokenizer) *)\n let input_ids =\n Nx.create Nx.int32 [| 4; 6 |]\n [|\n 101l;\n 1045l;\n 2293l;\n 2023l;\n 102l;\n 0l;\n (* \"I love this\" -> 1 *)\n 101l;\n 2307l;\n 3185l;\n 102l;\n 0l;\n 0l;\n (* \"great movie\" -> 1 *)\n 101l;\n 1045l;\n 5223l;\n 2023l;\n 102l;\n 0l;\n (* \"I hate this\" -> 0 *)\n 101l;\n 6659l;\n 2143l;\n 102l;\n 0l;\n 0l;\n (* \"terrible film\" -> 0 *)\n |]\n in\n let labels = Nx.create Nx.int32 [| 4 |] [| 1l; 1l; 0l; 0l |] in\n let attention_mask =\n Nx.create Nx.int32 [| 4; 6 |]\n [|\n 1l;\n 1l;\n 1l;\n 1l;\n 1l;\n 0l;\n 1l;\n 1l;\n 1l;\n 1l;\n 0l;\n 0l;\n 1l;\n 1l;\n 1l;\n 1l;\n 1l;\n 0l;\n 1l;\n 1l;\n 1l;\n 1l;\n 0l;\n 0l;\n |]\n in\n let ctx =\n Context.empty\n |> Context.set ~name:Attention.attention_mask_key (Ptree.P attention_mask)\n in\n\n (* --- Inference before training --- *)\n Printf.printf \"=== Before training ===\\n\";\n let logits_before =\n let y, _ = Layer.apply model vars ~training:false ~ctx input_ids in\n y\n in\n print_shape \"logits\" logits_before;\n\n (* --- Fine-tune --- *)\n Printf.printf \"\\n=== Training ===\\n%!\";\n let trainer =\n Train.make ~model\n ~optimizer:(Vega.adamw ~weight_decay:0.01 (Vega.Schedule.constant 2e-5))\n in\n let st = Train.make_state trainer vars in\n let st =\n Train.fit trainer st ~ctx\n ~report:(fun ~step ~loss _st ->\n Printf.printf \" step %2d loss %.4f\\n%!\" step loss)\n (Data.repeat 10\n (input_ids, fun logits -> Loss.cross_entropy_sparse logits labels))\n in\n\n (* --- Predictions after training --- *)\n Printf.printf \"\\n=== After training ===\\n\";\n let logits = Train.predict trainer st ~ctx input_ids in\n let sentences =\n [| \"I love this\"; \"great movie\"; \"I hate this\"; \"terrible film\" |]\n in\n for i = 0 to 3 do\n let row = Nx.slice [ I i ] logits in\n let v0 = Nx.item [ 0 ] row in\n let v1 = Nx.item [ 1 ] row in\n let pred = if v1 > v0 then \"positive\" else \"negative\" in\n let label = Int32.to_int (Nx.item [ i ] labels) in\n let expected = if label = 1 then \"positive\" else \"negative\" in\n Printf.printf \" %-20s pred=%-8s expected=%-8s %s\\n\"\n (Printf.sprintf \"\\\"%s\\\"\" sentences.(i))\n pred expected\n (if String.equal pred expected then \"OK\" else \"WRONG\")\n done\n"
},
{
"path": "packages/kaun/examples/03-bert/reference_hf_output.py",
"content": "#!/usr/bin/env python3\n\"\"\"Generate reference BERT outputs from HuggingFace transformers.\"\"\"\n\nfrom transformers import BertModel, BertTokenizer\nimport torch\n\n\ndef main():\n print(\"Generating HuggingFace BERT reference outputs\")\n print(\"=\" * 50)\n\n model = BertModel.from_pretrained(\"bert-base-uncased\")\n tokenizer = BertTokenizer.from_pretrained(\"bert-base-uncased\")\n\n text = \"Hello world\"\n print(f'\\nInput text: \"{text}\"')\n\n inputs = tokenizer(text, return_tensors=\"pt\")\n print(f\"Token IDs: {inputs['input_ids'][0].tolist()}\")\n\n with torch.no_grad():\n outputs = model(**inputs)\n\n last_hidden_state = outputs.last_hidden_state\n pooler_output = outputs.pooler_output\n\n print(f\"Output shape: {list(last_hidden_state.shape)}\")\n\n cls_token = last_hidden_state[0, 0]\n print(\"\\nCLS token (first 5 values):\")\n for i in range(5):\n print(f\" [{i}]: {cls_token[i].item():.6f}\")\n\n print(\"\\nPooler output (first 5 values):\")\n for i in range(5):\n print(f\" [{i}]: {pooler_output[0, i].item():.6f}\")\n\n print(\"\\n\" + \"=\" * 50)\n print(\"OCaml expected values:\")\n cls_vals = \"; \".join(f\"{cls_token[i].item():.6f}\" for i in range(5))\n pool_vals = \"; \".join(f\"{pooler_output[0, i].item():.6f}\" for i in range(5))\n print(f\"expected_cls = [| {cls_vals} |]\")\n print(f\"expected_pooler = [| {pool_vals} |]\")\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "packages/kaun/examples/04-gpt2/dune",
"content": "(executable\n (name main)\n (libraries nx nx.core rune kaun kaun.hf brot jsont))\n"
},
{
"path": "packages/kaun/examples/04-gpt2/gpt2.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Kaun\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\n\n(* Config *)\n\ntype config = {\n vocab_size : int;\n n_positions : int;\n n_embd : int;\n n_layer : int;\n n_head : int;\n n_inner : int;\n resid_pdrop : float;\n embd_pdrop : float;\n attn_pdrop : float;\n layer_norm_eps : float;\n}\n\nlet config ~vocab_size ~n_embd ~n_layer ~n_head ?(n_positions = 1024)\n ?(n_inner = 4 * n_embd) ?(resid_pdrop = 0.1) ?(embd_pdrop = 0.1)\n ?(attn_pdrop = 0.1) ?(layer_norm_eps = 1e-5) () =\n if n_embd mod n_head <> 0 then\n invalid_argf \"Gpt2.config: n_embd (%d) not divisible by n_head (%d)\" n_embd\n n_head;\n {\n vocab_size;\n n_positions;\n n_embd;\n n_layer;\n n_head;\n n_inner;\n resid_pdrop;\n embd_pdrop;\n attn_pdrop;\n layer_norm_eps;\n }\n\n(* Helpers *)\n\nlet fields ~ctx t = Ptree.Dict.fields_exn ~ctx t\nlet get fs ~name dtype = Ptree.Dict.get_tensor_exn fs ~name dtype\nlet find ~ctx key fs = Ptree.Dict.find_exn ~ctx key fs\n\n(* Causal self-attention with combined QKV *)\n\nlet causal_self_attention (type l) ~(cfg : config)\n ~(dtype : (float, l) Nx.dtype) ~training ~params (x : (float, l) Nx.t) :\n (float, l) Nx.t =\n let shape = Nx.shape x in\n let batch = shape.(0) in\n let seq = shape.(1) in\n let h = cfg.n_embd in\n let heads = cfg.n_head in\n let head_dim = h / heads in\n let fs = fields ~ctx:\"Gpt2.attention\" params in\n\n (* Combined QKV projection: [batch, seq, 3*h] *)\n let qkv_w = get fs ~name:\"qkv_weight\" dtype in\n let qkv_b = get fs ~name:\"qkv_bias\" dtype in\n let qkv = Nx.add (Nx.matmul x qkv_w) qkv_b in\n\n (* Split into Q, K, V *)\n let qkv_parts = Nx.split ~axis:(-1) 3 qkv in\n let q = List.nth qkv_parts 0 in\n let k = List.nth qkv_parts 1 in\n let v = List.nth qkv_parts 2 in\n\n let split_heads t =\n Nx.reshape [| batch; seq; heads; head_dim |] t\n |> Nx.transpose ~axes:[ 0; 2; 1; 3 ]\n in\n let q = split_heads q in\n let k = split_heads k in\n let v = split_heads v in\n\n let dropout_rate =\n if training && cfg.attn_pdrop > 0.0 then Some cfg.attn_pdrop else None\n in\n\n let attn =\n Kaun.Fn.dot_product_attention ~is_causal:true ?dropout_rate q k v\n in\n\n (* Merge heads *)\n let merged =\n Nx.transpose attn ~axes:[ 0; 2; 1; 3 ]\n |> Nx.contiguous\n |> Nx.reshape [| batch; seq; h |]\n in\n\n (* Output projection *)\n let o_w = get fs ~name:\"o_weight\" dtype in\n let o_b = get fs ~name:\"o_bias\" dtype in\n Nx.add (Nx.matmul merged o_w) o_b\n\n(* Transformer block (pre-norm) *)\n\nlet transformer_block (type l) ~(cfg : config) ~(dtype : (float, l) Nx.dtype)\n ~training ~params (x : (float, l) Nx.t) : (float, l) Nx.t =\n let fs = fields ~ctx:\"Gpt2.block\" params in\n\n (* Pre-norm attention *)\n let ln1_g = get fs ~name:\"ln1_gamma\" dtype in\n let ln1_b = get fs ~name:\"ln1_beta\" dtype in\n let x' =\n Kaun.Fn.layer_norm ~gamma:ln1_g ~beta:ln1_b ~epsilon:cfg.layer_norm_eps x\n in\n\n let attn_params = find ~ctx:\"Gpt2.block\" \"attention\" fs in\n let attn =\n causal_self_attention ~cfg ~dtype ~training ~params:attn_params x'\n in\n\n (* Residual dropout *)\n let attn =\n if training && cfg.resid_pdrop > 0.0 then\n Kaun.Fn.dropout ~rate:cfg.resid_pdrop attn\n else attn\n in\n let x = Nx.add x attn in\n\n (* Pre-norm FFN *)\n let ln2_g = get fs ~name:\"ln2_gamma\" dtype in\n let ln2_b = get fs ~name:\"ln2_beta\" dtype in\n let x' =\n Kaun.Fn.layer_norm ~gamma:ln2_g ~beta:ln2_b ~epsilon:cfg.layer_norm_eps x\n in\n\n let ffn_up_w = get fs ~name:\"ffn_up_weight\" dtype in\n let ffn_up_b = get fs ~name:\"ffn_up_bias\" dtype in\n let ffn_down_w = get fs ~name:\"ffn_down_weight\" dtype in\n let ffn_down_b = get fs ~name:\"ffn_down_bias\" dtype in\n\n let y =\n Nx.add (Nx.matmul x' ffn_up_w) ffn_up_b |> Kaun.Activation.gelu_approx\n in\n let y = Nx.add (Nx.matmul y ffn_down_w) ffn_down_b in\n\n (* Residual dropout *)\n let y =\n if training && cfg.resid_pdrop > 0.0 then\n Kaun.Fn.dropout ~rate:cfg.resid_pdrop y\n else y\n in\n Nx.add x y\n\n(* Forward: embeddings + transformer stack + final layer norm *)\n\nlet decode (type l in_elt) ~(cfg : config) ~params\n ~(dtype : (float, l) Nx.dtype) ~training (input_ids : (int32, in_elt) Nx.t)\n : (float, l) Nx.t =\n let input_ids = Nx.cast Nx.int32 input_ids in\n let shape = Nx.shape input_ids in\n let batch = shape.(0) in\n let seq = shape.(1) in\n\n if seq > cfg.n_positions then\n invalid_argf \"Gpt2.decode: seq_len=%d exceeds n_positions=%d\" seq\n cfg.n_positions;\n\n (* Params *)\n let root = fields ~ctx:\"Gpt2.decode\" params in\n\n let wte = get root ~name:\"wte\" dtype in\n let wpe = get root ~name:\"wpe\" dtype in\n let layers_t = find ~ctx:\"Gpt2.decode\" \"layers\" root in\n\n (* Embedding lookup: token + position *)\n let position_ids =\n Nx.arange_f Nx.float32 0.0 (float_of_int seq) 1.0\n |> Nx.cast Nx.int32\n |> Nx.reshape [| 1; seq |]\n |> Nx.broadcast_to [| batch; seq |]\n |> Nx.contiguous\n in\n let tok = Kaun.Fn.embedding ~scale:false ~embedding:wte input_ids in\n let pos = Kaun.Fn.embedding ~scale:false ~embedding:wpe position_ids in\n let x = Nx.add tok pos in\n\n (* Embedding dropout *)\n let x =\n if training && cfg.embd_pdrop > 0.0 then\n Kaun.Fn.dropout ~rate:cfg.embd_pdrop x\n else x\n in\n\n (* Transformer stack *)\n let blocks = Ptree.List.items_exn ~ctx:\"Gpt2.decode.layers\" layers_t in\n let x =\n List.fold_left\n (fun h block_params ->\n transformer_block ~cfg ~dtype ~training ~params:block_params h)\n x blocks\n in\n\n (* Final layer norm *)\n let ln_f_g = get root ~name:\"ln_f_gamma\" dtype in\n let ln_f_b = get root ~name:\"ln_f_beta\" dtype in\n Kaun.Fn.layer_norm ~gamma:ln_f_g ~beta:ln_f_b ~epsilon:cfg.layer_norm_eps x\n\n(* Parameter initialization *)\n\nlet init_block_params ~dtype ~n_embd ~n_inner =\n let w = Init.normal ~stddev:0.02 () in\n let zeros n = Nx.zeros dtype [| n |] in\n let ones n = Nx.ones dtype [| n |] in\n let attn_params =\n Ptree.dict\n [\n (\"qkv_weight\", Ptree.tensor (w.f [| n_embd; 3 * n_embd |] dtype));\n (\"qkv_bias\", Ptree.tensor (zeros (3 * n_embd)));\n (\"o_weight\", Ptree.tensor (w.f [| n_embd; n_embd |] dtype));\n (\"o_bias\", Ptree.tensor (zeros n_embd));\n ]\n in\n Ptree.dict\n [\n (\"attention\", attn_params);\n (\"ln1_gamma\", Ptree.tensor (ones n_embd));\n (\"ln1_beta\", Ptree.tensor (zeros n_embd));\n (\"ffn_up_weight\", Ptree.tensor (w.f [| n_embd; n_inner |] dtype));\n (\"ffn_up_bias\", Ptree.tensor (zeros n_inner));\n (\"ffn_down_weight\", Ptree.tensor (w.f [| n_inner; n_embd |] dtype));\n (\"ffn_down_bias\", Ptree.tensor (zeros n_embd));\n (\"ln2_gamma\", Ptree.tensor (ones n_embd));\n (\"ln2_beta\", Ptree.tensor (zeros n_embd));\n ]\n\nlet init_decoder_params ~cfg ~dtype =\n let h = cfg.n_embd in\n let w = Init.normal ~stddev:0.02 () in\n let wte = w.f [| cfg.vocab_size; h |] dtype in\n let wpe = w.f [| cfg.n_positions; h |] dtype in\n let blocks =\n List.init cfg.n_layer (fun _ ->\n init_block_params ~dtype ~n_embd:h ~n_inner:cfg.n_inner)\n in\n Ptree.dict\n [\n (\"wte\", Ptree.tensor wte);\n (\"wpe\", Ptree.tensor wpe);\n (\"layers\", Ptree.list blocks);\n (\"ln_f_gamma\", Ptree.tensor (Nx.ones dtype [| h |]));\n (\"ln_f_beta\", Ptree.tensor (Nx.zeros dtype [| h |]));\n ]\n\n(* Layers *)\n\nlet decoder (cfg : config) () : (int32, float) Layer.t =\n {\n Layer.init =\n (fun ~dtype ->\n Layer.make_vars\n ~params:(init_decoder_params ~cfg ~dtype)\n ~state:Ptree.empty ~dtype);\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (state, ctx);\n let y = decode ~cfg ~params ~dtype ~training x in\n (y, Ptree.empty));\n }\n\nlet for_causal_lm (cfg : config) () : (int32, float) Layer.t =\n {\n Layer.init =\n (fun ~dtype ->\n Layer.make_vars\n ~params:(init_decoder_params ~cfg ~dtype)\n ~state:Ptree.empty ~dtype);\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (state, ctx);\n let hidden = decode ~cfg ~params ~dtype ~training x in\n (* Tied LM head: logits = hidden @ wte^T *)\n let root = fields ~ctx:\"Gpt2.lm_head\" params in\n let wte = get root ~name:\"wte\" dtype in\n let logits = Nx.matmul hidden (Nx.transpose wte ~axes:[ 1; 0 ]) in\n (logits, Ptree.empty));\n }\n\n(* JSON config parsing *)\n\nlet json_mem name = function\n | Jsont.Object (mems, _) -> (\n match Jsont.Json.find_mem name mems with\n | Some (_, v) -> v\n | None -> Jsont.Null ((), Jsont.Meta.none))\n | _ -> Jsont.Null ((), Jsont.Meta.none)\n\nlet json_to_int = function\n | Jsont.Number (f, _) -> int_of_float f\n | _ -> failwith \"expected int\"\n\nlet json_to_int_option = function\n | Jsont.Number (f, _) -> Some (int_of_float f)\n | _ -> None\n\nlet json_to_float_option = function Jsont.Number (f, _) -> Some f | _ -> None\n\nlet parse_config json =\n let n_embd = json |> json_mem \"n_embd\" |> json_to_int in\n config\n ~vocab_size:(json |> json_mem \"vocab_size\" |> json_to_int)\n ~n_embd\n ~n_layer:(json |> json_mem \"n_layer\" |> json_to_int)\n ~n_head:(json |> json_mem \"n_head\" |> json_to_int)\n ?n_positions:(json |> json_mem \"n_positions\" |> json_to_int_option)\n ?n_inner:(json |> json_mem \"n_inner\" |> json_to_int_option)\n ?resid_pdrop:(json |> json_mem \"resid_pdrop\" |> json_to_float_option)\n ?embd_pdrop:(json |> json_mem \"embd_pdrop\" |> json_to_float_option)\n ?attn_pdrop:(json |> json_mem \"attn_pdrop\" |> json_to_float_option)\n ?layer_norm_eps:\n (json |> json_mem \"layer_norm_epsilon\" |> json_to_float_option)\n ()\n\n(* HuggingFace weight mapping *)\n\nlet cast_tensor dtype (Ptree.P t) = Ptree.P (Nx.cast dtype t)\n\nlet map_hf_weights ~cfg ~dtype hf_weights =\n let tbl = Hashtbl.create (List.length hf_weights) in\n List.iter (fun (name, tensor) -> Hashtbl.add tbl name tensor) hf_weights;\n let hf name =\n match Hashtbl.find_opt tbl name with\n | Some t -> cast_tensor dtype t\n | None -> invalid_argf \"from_pretrained: missing HF weight %S\" name\n in\n (* GPT-2 stores weights as [in, out] — NO transpose needed *)\n let hf_t name = Ptree.Tensor (hf name) in\n let layer i =\n let p s = Printf.sprintf \"h.%d.%s\" i s in\n Ptree.dict\n [\n ( \"attention\",\n Ptree.dict\n [\n (\"qkv_weight\", hf_t (p \"attn.c_attn.weight\"));\n (\"qkv_bias\", hf_t (p \"attn.c_attn.bias\"));\n (\"o_weight\", hf_t (p \"attn.c_proj.weight\"));\n (\"o_bias\", hf_t (p \"attn.c_proj.bias\"));\n ] );\n (\"ln1_gamma\", hf_t (p \"ln_1.weight\"));\n (\"ln1_beta\", hf_t (p \"ln_1.bias\"));\n (\"ffn_up_weight\", hf_t (p \"mlp.c_fc.weight\"));\n (\"ffn_up_bias\", hf_t (p \"mlp.c_fc.bias\"));\n (\"ffn_down_weight\", hf_t (p \"mlp.c_proj.weight\"));\n (\"ffn_down_bias\", hf_t (p \"mlp.c_proj.bias\"));\n (\"ln2_gamma\", hf_t (p \"ln_2.weight\"));\n (\"ln2_beta\", hf_t (p \"ln_2.bias\"));\n ]\n in\n Ptree.dict\n [\n (\"wte\", hf_t \"wte.weight\");\n (\"wpe\", hf_t \"wpe.weight\");\n (\"layers\", Ptree.list (List.init cfg.n_layer layer));\n (\"ln_f_gamma\", hf_t \"ln_f.weight\");\n (\"ln_f_beta\", hf_t \"ln_f.bias\");\n ]\n\n(* Pretrained loading *)\n\nlet from_pretrained ?(model_id = \"gpt2\") () =\n let json = Kaun_hf.load_config ~model_id () in\n let cfg = parse_config json in\n let hf_weights = Kaun_hf.load_weights ~model_id () in\n let params = map_hf_weights ~cfg ~dtype:Nx.float32 hf_weights in\n (cfg, params)\n"
},
{
"path": "packages/kaun/examples/04-gpt2/gpt2.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** GPT-2 decoder and language model head.\n\n GPT-2 inputs are passed as int32 [input_ids]:\n\n {[\n Layer.apply model vars ~training:false input_ids\n ]}\n\n Position ids are computed automatically from the sequence length. *)\n\nopen Kaun\n\n(** {1:config Configuration} *)\n\ntype config = {\n vocab_size : int;\n n_positions : int;\n n_embd : int;\n n_layer : int;\n n_head : int;\n n_inner : int;\n resid_pdrop : float;\n embd_pdrop : float;\n attn_pdrop : float;\n layer_norm_eps : float;\n}\n(** The type for GPT-2 configurations. *)\n\nval config :\n vocab_size:int ->\n n_embd:int ->\n n_layer:int ->\n n_head:int ->\n ?n_positions:int ->\n ?n_inner:int ->\n ?resid_pdrop:float ->\n ?embd_pdrop:float ->\n ?attn_pdrop:float ->\n ?layer_norm_eps:float ->\n unit ->\n config\n(** [config ~vocab_size ~n_embd ~n_layer ~n_head ()] is a GPT-2 configuration.\n\n [n_positions] defaults to [1024]. [n_inner] defaults to [4 * n_embd].\n Dropout rates default to [0.1]. [layer_norm_eps] defaults to [1e-5].\n\n Raises [Invalid_argument] if [n_embd] is not divisible by [n_head]. *)\n\n(** {1:layers Layers} *)\n\nval decoder : config -> unit -> (int32, float) Layer.t\n(** [decoder cfg ()] is the GPT-2 transformer decoder.\n\n Input: int32 [input_ids] of shape [[batch; seq]]. Output: float hidden\n states of shape [[batch; seq; n_embd]]. *)\n\nval for_causal_lm : config -> unit -> (int32, float) Layer.t\n(** [for_causal_lm cfg ()] is decoder + tied LM head.\n\n Output: logits [[batch; seq; vocab_size]]. Word embeddings are tied with the\n LM head projection. *)\n\n(** {1:pretrained Pretrained loading} *)\n\nval from_pretrained : ?model_id:string -> unit -> config * Ptree.t\n(** [from_pretrained ?model_id ()] downloads [model_id] from HuggingFace and\n returns [(cfg, decoder_params)].\n\n [decoder_params] is ready for {!decoder} or {!for_causal_lm} (the LM head\n reuses the word embedding weights).\n\n [model_id] defaults to [\"gpt2\"]. *)\n"
},
{
"path": "packages/kaun/examples/04-gpt2/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Autoregressive text generation with pretrained GPT-2.\n\n Downloads gpt2 from HuggingFace (~548MB on first run) and generates text\n continuations from several prompts using greedy decoding. *)\n\nopen Kaun\n\n(* Tokenizer *)\n\nlet load_tokenizer model_id =\n let vocab = Kaun_hf.download_file ~model_id ~filename:\"vocab.json\" () in\n let merges = Kaun_hf.download_file ~model_id ~filename:\"merges.txt\" () in\n Brot.from_model_file ~vocab ~merges\n ~pre:\n (Brot.Pre_tokenizer.byte_level ~add_prefix_space:false ~use_regex:true ())\n ~decoder:(Brot.Decoder.byte_level ())\n ()\n\nlet encode tokenizer text =\n Array.map Int32.of_int (Brot.encode_ids tokenizer text)\n\nlet decode tokenizer ids = Brot.decode tokenizer (Array.map Int32.to_int ids)\n\n(* Greedy decode: at each step pick the highest-probability next token. *)\n\nlet generate model vars ~max_tokens prompt =\n let tokens = ref (Array.to_list prompt) in\n for _ = 1 to max_tokens do\n let ids = Array.of_list !tokens in\n let n = Array.length ids in\n let input = Nx.create Nx.int32 [| 1; n |] ids in\n let logits, _ = Layer.apply model vars ~training:false input in\n let last = Nx.slice [ I 0; I (n - 1) ] logits in\n let next : int32 = Nx.item [] (Nx.argmax ~axis:0 last) in\n tokens := !tokens @ [ next ]\n done;\n Array.of_list !tokens\n\n(* Show the model's top-k predictions at a given position. *)\n\nlet print_top_k ~k model vars input_ids ~pos =\n let logits, _ = Layer.apply model vars ~training:false input_ids in\n let row = Nx.slice [ I 0; I pos ] logits in\n let sorted = Nx.argsort ~descending:true ~axis:0 row in\n let probs = Nx.softmax ~axes:[ 0 ] row in\n for i = 0 to k - 1 do\n let idx = Int32.to_int (Nx.item [ i ] sorted) in\n let prob : float = Nx.item [ idx ] probs in\n Printf.printf \" #%d token %-6d p=%.4f\\n\" (i + 1) idx prob\n done\n\nlet () =\n let model_id = \"gpt2\" in\n let dtype = Nx.float32 in\n\n (* Load tokenizer and model *)\n Printf.printf \"Loading %s...\\n%!\" model_id;\n let tokenizer = load_tokenizer model_id in\n let cfg, params = Gpt2.from_pretrained ~model_id () in\n Printf.printf \" vocab=%d n_embd=%d layers=%d heads=%d\\n\\n\" cfg.vocab_size\n cfg.n_embd cfg.n_layer cfg.n_head;\n\n let model = Gpt2.for_causal_lm cfg () in\n let vars = Layer.make_vars ~params ~state:Ptree.empty ~dtype in\n\n (* --- What does the model predict after \"Hello world\"? --- *)\n Printf.printf \"=== Next-token predictions ===\\n\";\n Printf.printf \" Prompt: \\\"Hello world\\\"\\n\";\n Printf.printf \" Top 5 continuations:\\n\";\n let hello_ids = encode tokenizer \"Hello world\" in\n let hello = Nx.create Nx.int32 [| 1; Array.length hello_ids |] hello_ids in\n print_top_k ~k:5 model vars hello ~pos:(Array.length hello_ids - 1);\n\n (* --- Greedy generation from several prompts --- *)\n Printf.printf \"\\n=== Greedy generation (30 tokens each) ===\\n\\n\";\n let prompts =\n [ \"The meaning of life is\"; \"Once upon a time\"; \"The quick brown fox\" ]\n in\n List.iter\n (fun text ->\n let prompt = encode tokenizer text in\n let generated = generate model vars ~max_tokens:30 prompt in\n let continuation =\n Array.sub generated (Array.length prompt)\n (Array.length generated - Array.length prompt)\n in\n Printf.printf \" \\\"%s\\\" ->\\n\" text;\n Printf.printf \" %s\\n\\n\" (decode tokenizer continuation))\n prompts\n"
},
{
"path": "packages/kaun/examples/04-gpt2/reference_hf_output.py",
"content": "#!/usr/bin/env python3\n\"\"\"Get reference GPT-2 output from transformers for comparison.\"\"\"\n\nimport torch\nfrom transformers import GPT2LMHeadModel, GPT2Tokenizer\n\nprint(\"Loading GPT-2 model and tokenizer...\")\ntokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\nmodel = GPT2LMHeadModel.from_pretrained(\"gpt2\")\nmodel.eval()\n\ntext = \"Hello world\"\nprint(f\"Test input: {repr(text)}\")\n\ninputs = tokenizer(text, return_tensors=\"pt\")\nprint(f\"Input IDs: {inputs.input_ids.tolist()[0]}\")\n\nwith torch.no_grad():\n outputs = model(**inputs)\n logits = outputs.logits\n\nprint(f\"\\nLogits shape: {list(logits.shape)}\")\nprint(f\"First 10 logit values at position 0: {logits[0, 0, :10].tolist()}\")\n\n# Top 5 predictions for next token (after \"world\")\nlast_logits = logits[0, -1, :]\nprobs = torch.softmax(last_logits, dim=-1)\ntop_probs, top_indices = torch.topk(probs, 5)\n\nprint(\"\\nTop 5 predicted next tokens:\")\nfor idx, prob in zip(top_indices.tolist(), top_probs.tolist()):\n token = tokenizer.decode([idx])\n print(f\" Token {idx} ({repr(token)}): {prob:.4f}\")\n\n# Greedy generation\ngenerated = model.generate(\n inputs.input_ids, max_new_tokens=20, do_sample=False\n)\nprint(f\"\\nGreedy generation: {tokenizer.decode(generated[0])}\")\nprint(f\"Token IDs: {generated[0].tolist()}\")\n"
},
{
"path": "packages/kaun/lib/activation.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Re-exports *)\n\nlet relu x = Nx.relu x\nlet sigmoid x = Nx.sigmoid x\nlet tanh x = Nx.tanh x\n\n(* Activations *)\n\nlet relu6 x =\n let zero = Nx.scalar_like x 0.0 in\n let six = Nx.scalar_like x 6.0 in\n Nx.minimum (Nx.maximum x zero) six\n\nlet hard_sigmoid ?(alpha = 1.0 /. 6.0) ?(beta = 0.5) x =\n let linear =\n Nx.add (Nx.mul (Nx.scalar_like x alpha) x) (Nx.scalar_like x beta)\n in\n let zero = Nx.scalar_like x 0. in\n let one = Nx.scalar_like x 1. in\n Nx.minimum one (Nx.maximum zero linear)\n\nlet softplus x = Nx.log (Nx.add (Nx.scalar_like x 1.) (Nx.exp x))\nlet silu x = Nx.mul x (Nx.sigmoid x)\nlet swish x = silu x\nlet hard_silu x = Nx.mul x (hard_sigmoid x)\nlet hard_swish x = hard_silu x\n\nlet prelu ~alpha x =\n let zero = Nx.zeros_like x in\n Nx.add (Nx.maximum zero x) (Nx.mul alpha (Nx.minimum zero x))\n\nlet log_sigmoid x =\n (* Numerically stable: branch on sign to avoid overflow *)\n let zero = Nx.scalar_like x 0.0 in\n let one = Nx.scalar_like x 1.0 in\n let is_positive = Nx.greater x zero in\n let branch_pos = Nx.neg (Nx.log (Nx.add one (Nx.exp (Nx.neg x)))) in\n let branch_neg = Nx.sub x (Nx.log (Nx.add one (Nx.exp x))) in\n Nx.where is_positive branch_pos branch_neg\n\nlet leaky_relu ?(negative_slope = 0.01) x =\n Nx.maximum x (Nx.mul (Nx.scalar_like x negative_slope) x)\n\nlet hard_tanh x =\n let one = Nx.scalar_like x 1. in\n let neg_one = Nx.scalar_like x (-1.0) in\n Nx.maximum neg_one (Nx.minimum x one)\n\nlet elu ?(alpha = 1.0) x =\n let zero = Nx.scalar_like x 0.0 in\n let one = Nx.scalar_like x 1. in\n let alpha_s = Nx.scalar_like x alpha in\n let exp_minus_one = Nx.sub (Nx.exp x) one in\n Nx.add (Nx.maximum x zero) (Nx.mul alpha_s (Nx.minimum zero exp_minus_one))\n\nlet selu x =\n let alpha = 1.6732632423543772848170429916717 in\n let lambda = 1.0507009873554804934193349852946 in\n Nx.mul (Nx.scalar_like x lambda) (elu ~alpha x)\n\nlet celu ?(alpha = 1.0) x =\n let zero = Nx.zeros_like x in\n let alpha_s = Nx.scalar_like x alpha in\n let one = Nx.scalar_like x 1. in\n let neg_term =\n Nx.mul alpha_s (Nx.sub (Nx.exp (Nx.div (Nx.minimum zero x) alpha_s)) one)\n in\n Nx.add (Nx.maximum zero x) neg_term\n\nlet squareplus ?(b = 4.0) x =\n let half = Nx.scalar_like x 0.5 in\n let inside = Nx.add (Nx.square x) (Nx.scalar_like x b) in\n Nx.mul half (Nx.add x (Nx.sqrt inside))\n\nlet glu ?(axis = -1) x =\n match Nx.split ~axis 2 x with\n | [ left; right ] -> Nx.mul left (Nx.sigmoid right)\n | _ -> invalid_arg \"Activation.glu: split did not produce two partitions\"\n\nlet sparse_plus x =\n let zero = Nx.zeros_like x in\n let one = Nx.scalar_like x 1. in\n let neg_one = Nx.scalar_like x (-1.) in\n let quadratic = Nx.mul (Nx.scalar_like x 0.25) (Nx.square (Nx.add x one)) in\n let res = Nx.where (Nx.greater_equal x one) x quadratic in\n Nx.where (Nx.less_equal x neg_one) zero res\n\nlet sparse_sigmoid x =\n let zero = Nx.zeros_like x in\n let one = Nx.scalar_like x 1. in\n let neg_one = Nx.scalar_like x (-1.) in\n let half = Nx.scalar_like x 0.5 in\n let linear = Nx.mul half (Nx.add x one) in\n let res = Nx.where (Nx.greater_equal x one) one linear in\n Nx.where (Nx.less_equal x neg_one) zero res\n\nlet gelu_approx x =\n let one = Nx.scalar_like x 1.0 in\n let half = Nx.scalar_like x 0.5 in\n let sqrt2_pi = Nx.scalar_like x 0.7978845608 in\n let coeff = Nx.scalar_like x 0.044715 in\n let x2 = Nx.mul x x in\n let inner = Nx.add one (Nx.mul coeff x2) in\n let arg = Nx.mul (Nx.mul x sqrt2_pi) inner in\n Nx.mul half (Nx.mul x (Nx.add one (Nx.tanh arg)))\n\nlet gelu x =\n let half = Nx.scalar_like x 0.5 in\n let one = Nx.scalar_like x 1.0 in\n let sqrt2 = Nx.scalar_like x 1.4142135623730951 in\n Nx.mul (Nx.mul half x) (Nx.add one (Nx.erf (Nx.div x sqrt2)))\n\nlet softsign x =\n let one = Nx.scalar_like x 1.0 in\n Nx.div x (Nx.add one (Nx.abs x))\n\nlet mish x = Nx.mul x (Nx.tanh (softplus x))\n"
},
{
"path": "packages/kaun/lib/activation.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Activation functions for neural networks.\n\n All functions are differentiable through Rune's autodiff. The standard\n activations {!relu}, {!sigmoid} and {!tanh} are re-exported from {!Nx} for\n convenience. *)\n\n(** {1:standard Standard activations} *)\n\nval relu : ('a, 'b) Nx.t -> ('a, 'b) Nx.t\n(** [relu x] is [max(x, 0)]. *)\n\nval sigmoid : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [sigmoid x] is [1 / (1 + exp(-x))]. *)\n\nval tanh : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [tanh x] is the hyperbolic tangent of [x]. *)\n\nval relu6 : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [relu6 x] is [min(max(x, 0), 6)]. *)\n\nval leaky_relu : ?negative_slope:float -> (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [leaky_relu x] is [max(x, negative_slope * x)].\n\n [negative_slope] defaults to [0.01]. *)\n\nval hard_tanh : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [hard_tanh x] is [max(-1, min(1, x))]. *)\n\nval hard_sigmoid :\n ?alpha:float -> ?beta:float -> (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [hard_sigmoid x] is [min(1, max(0, alpha * x + beta))].\n\n [alpha] defaults to [1/6]. [beta] defaults to [0.5]. *)\n\nval prelu : alpha:(float, 'b) Nx.t -> (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [prelu ~alpha x] is [max(0, x) + alpha * min(0, x)].\n\n [alpha] is a learnable tensor, broadcast against [x]. *)\n\n(** {1:exponential Exponential family} *)\n\nval elu : ?alpha:float -> (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [elu x] is [x] when [x >= 0] and [alpha * (exp(x) - 1)] otherwise.\n\n [alpha] defaults to [1.0]. *)\n\nval selu : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [selu x] is [lambda * elu(x, alpha)] with self-normalizing constants. *)\n\nval celu : ?alpha:float -> (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [celu x] is [max(0, x) + min(0, alpha * (exp(x/alpha) - 1))].\n\n [alpha] defaults to [1.0]. *)\n\n(** {1:smooth Smooth activations} *)\n\nval gelu : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [gelu x] is [0.5 * x * (1 + erf(x / sqrt(2)))]. *)\n\nval gelu_approx : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [gelu_approx x] is the tanh-based approximation of {!gelu}. *)\n\nval silu : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [silu x] is [x * sigmoid(x)] (also known as Swish). *)\n\nval swish : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [swish x] is {!silu}. *)\n\nval hard_silu : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [hard_silu x] is [x * hard_sigmoid(x)]. *)\n\nval hard_swish : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [hard_swish x] is {!hard_silu}. *)\n\nval mish : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [mish x] is [x * tanh(softplus(x))]. *)\n\nval softplus : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [softplus x] is [log(1 + exp(x))]. *)\n\nval softsign : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [softsign x] is [x / (abs(x) + 1)]. *)\n\nval squareplus : ?b:float -> (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [squareplus x] is [0.5 * (x + sqrt(x^2 + b))].\n\n [b] defaults to [4.0]. *)\n\nval log_sigmoid : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [log_sigmoid x] is [log(sigmoid(x))], computed in a numerically stable way\n by branching on the sign of [x]. *)\n\n(** {1:gating Gating} *)\n\nval glu : ?axis:int -> (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [glu x] splits [x] in half along [axis] and returns [left * sigmoid(right)].\n\n [axis] defaults to [-1].\n\n Raises [Invalid_argument] if the split does not produce two partitions. *)\n\n(** {1:sparse Sparse activations} *)\n\nval sparse_plus : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [sparse_plus x] is [x] when [x >= 1], [0] when [x <= -1], and\n [0.25 * (x + 1)^2] otherwise. *)\n\nval sparse_sigmoid : (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [sparse_sigmoid x] is [1] when [x >= 1], [0] when [x <= -1], and\n [0.5 * (x + 1)] otherwise. *)\n"
},
{
"path": "packages/kaun/lib/attention.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\n\nmodule Dtype = Nx_core.Dtype\n\nlet require_same_float_dtype (type p in_elt) ~ctx\n (expected : (float, p) Nx.dtype) (x : (float, in_elt) Nx.t) :\n (float, p) Nx.t =\n match Dtype.equal_witness expected (Nx.dtype x) with\n | Some Type.Equal -> (x : (float, p) Nx.t)\n | None ->\n invalid_argf \"%s: input dtype %s does not match model dtype %s\" ctx\n (Dtype.to_string (Nx.dtype x))\n (Dtype.to_string expected)\n\nlet normalize_axis ~ctx ~ndim axis =\n let normalized = if axis < 0 then ndim + axis else axis in\n if normalized < 0 || normalized >= ndim then\n invalid_argf \"%s: axis %d out of bounds for rank %d\" ctx axis ndim;\n normalized\n\n(* Rotary position embeddings *)\n\nlet rope ?(theta = 10000.0) ?(seq_dim = -2) x =\n let ctx = \"Attention.rope\" in\n let shape = Nx.shape x in\n let ndim = Array.length shape in\n if ndim < 2 then invalid_argf \"%s: expected rank >= 2, got rank %d\" ctx ndim;\n let seq_axis = normalize_axis ~ctx ~ndim seq_dim in\n if seq_axis = ndim - 1 then\n invalid_argf\n \"%s: seq_dim points to the last axis; last axis is reserved for head_dim\"\n ctx;\n let seq_len = shape.(seq_axis) in\n let head_dim = shape.(ndim - 1) in\n if head_dim mod 2 <> 0 then\n invalid_argf \"%s: head_dim must be even, got %d\" ctx head_dim;\n let half = head_dim / 2 in\n let dtype = Nx.dtype x in\n let inv_freq =\n let exponents = Nx.arange_f dtype 0.0 (float_of_int head_dim) 2.0 in\n let normalized =\n Nx.div exponents (Nx.scalar dtype (float_of_int head_dim))\n in\n Nx.pow (Nx.scalar dtype theta) (Nx.neg normalized)\n in\n let positions = Nx.arange_f dtype 0.0 (float_of_int seq_len) 1.0 in\n let angles =\n Nx.matmul\n (Nx.reshape [| seq_len; 1 |] positions)\n (Nx.reshape [| 1; half |] inv_freq)\n in\n let broadcast_shape = Array.make ndim 1 in\n broadcast_shape.(seq_axis) <- seq_len;\n broadcast_shape.(ndim - 1) <- half;\n let cos_angles = Nx.reshape broadcast_shape (Nx.cos angles) in\n let sin_angles = Nx.reshape broadcast_shape (Nx.sin angles) in\n let last_axis_slice start stop =\n let slices = Array.make ndim Nx.A in\n slices.(ndim - 1) <- Nx.R (start, stop);\n Array.to_list slices\n in\n let x1 = Nx.slice (last_axis_slice 0 half) x in\n let x2 = Nx.slice (last_axis_slice half head_dim) x in\n let r1 = Nx.sub (Nx.mul x1 cos_angles) (Nx.mul x2 sin_angles) in\n let r2 = Nx.add (Nx.mul x1 sin_angles) (Nx.mul x2 cos_angles) in\n Nx.concatenate ~axis:(-1) [ r1; r2 ]\n\n(* Multi-head self-attention *)\n\nlet attention_mask_key = \"attention_mask\"\nlet apply_rope ~theta t = rope ~theta t\n\nlet multi_head_attention ~embed_dim ~num_heads ?(num_kv_heads = num_heads)\n ?(dropout = 0.0) ?(is_causal = false) ?(rope = false)\n ?(rope_theta = 10000.0) () =\n let use_rope = rope in\n let head_dim = embed_dim / num_heads in\n if head_dim * num_heads <> embed_dim then\n invalid_argf\n \"Attention.multi_head_attention: embed_dim (%d) not divisible by \\\n num_heads (%d)\"\n embed_dim num_heads;\n if num_heads mod num_kv_heads <> 0 then\n invalid_argf\n \"Attention.multi_head_attention: num_heads (%d) not divisible by \\\n num_kv_heads (%d)\"\n num_heads num_kv_heads;\n if dropout < 0.0 || dropout >= 1.0 then\n invalid_argf\n \"Attention.multi_head_attention: expected 0.0 <= dropout < 1.0, got %g\"\n dropout;\n let weight_init = Init.glorot_uniform () in\n {\n Layer.init =\n (fun ~dtype ->\n let q_proj =\n weight_init.f [| embed_dim; num_heads * head_dim |] dtype\n in\n let k_proj =\n weight_init.f [| embed_dim; num_kv_heads * head_dim |] dtype\n in\n let v_proj =\n weight_init.f [| embed_dim; num_kv_heads * head_dim |] dtype\n in\n let out_proj =\n weight_init.f [| num_heads * head_dim; embed_dim |] dtype\n in\n Layer.make_vars\n ~params:\n (Ptree.dict\n [\n (\"q_proj\", Ptree.tensor q_proj);\n (\"k_proj\", Ptree.tensor k_proj);\n (\"v_proj\", Ptree.tensor v_proj);\n (\"out_proj\", Ptree.tensor out_proj);\n ])\n ~state:(Ptree.list []) ~dtype);\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n let x =\n require_same_float_dtype ~ctx:\"Attention.multi_head_attention\" dtype x\n in\n let shape = Nx.shape x in\n let batch = shape.(0) in\n let seq_len = shape.(1) in\n let fields =\n Ptree.Dict.fields_exn ~ctx:\"Attention.multi_head_attention.params\"\n params\n in\n let get name = Ptree.Dict.get_tensor_exn fields ~name dtype in\n let q_proj = get \"q_proj\" in\n let k_proj = get \"k_proj\" in\n let v_proj = get \"v_proj\" in\n let out_proj = get \"out_proj\" in\n let q = Nx.matmul x q_proj in\n let k = Nx.matmul x k_proj in\n let v = Nx.matmul x v_proj in\n let reshape_heads t heads =\n let t = Nx.reshape [| batch; seq_len; heads; head_dim |] t in\n Nx.transpose t ~axes:[ 0; 2; 1; 3 ]\n in\n let q = reshape_heads q num_heads in\n let k = reshape_heads k num_kv_heads in\n let v = reshape_heads v num_kv_heads in\n let repeat_kv t =\n if num_kv_heads < num_heads then\n let repetition = num_heads / num_kv_heads in\n let shape = Nx.shape t in\n let expanded = Nx.expand_dims [ 2 ] t in\n let target =\n [| shape.(0); shape.(1); repetition; shape.(2); shape.(3) |]\n in\n Nx.broadcast_to target expanded\n |> Nx.contiguous\n |> Nx.reshape [| shape.(0); num_heads; shape.(2); shape.(3) |]\n else t\n in\n let k = repeat_kv k in\n let v = repeat_kv v in\n let q, k =\n if use_rope then\n (apply_rope ~theta:rope_theta q, apply_rope ~theta:rope_theta k)\n else (q, k)\n in\n let dropout_rate =\n if training && dropout > 0.0 then Some dropout else None\n in\n (* Read attention mask from context if present. Accepts [batch; seq_k]\n int32 (0/1) or bool, reshapes to [batch; 1; 1; seq_k] for\n broadcasting over heads and queries. *)\n let attention_mask =\n match ctx with\n | None -> None\n | Some ctx -> (\n match Context.find ctx ~name:attention_mask_key with\n | None -> None\n | Some tensor ->\n let bool_mask =\n match Ptree.Tensor.to_typed Nx.bool tensor with\n | Some m -> m\n | None ->\n (* int/float mask: cast to int32, nonzero = true *)\n let int_mask =\n match Ptree.Tensor.to_typed Nx.int32 tensor with\n | Some m -> m\n | None ->\n let (Ptree.P raw) = tensor in\n Nx.cast Nx.int32 raw\n in\n Nx.not_equal int_mask\n (Nx.zeros Nx.int32 (Nx.shape int_mask))\n in\n let mask_shape = Nx.shape bool_mask in\n let ndim = Array.length mask_shape in\n let reshaped =\n if ndim = 2 then\n Nx.reshape\n [| mask_shape.(0); 1; 1; mask_shape.(1) |]\n bool_mask\n else bool_mask\n in\n Some reshaped)\n in\n let attn =\n Fn.dot_product_attention ?attention_mask ?dropout_rate ~is_causal q k\n v\n in\n let merged =\n Nx.transpose attn ~axes:[ 0; 2; 1; 3 ]\n |> Nx.contiguous\n |> Nx.reshape [| batch; seq_len; embed_dim |]\n in\n let output = Nx.matmul merged out_proj in\n (output, state));\n }\n"
},
{
"path": "packages/kaun/lib/attention.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Multi-head self-attention.\n\n Provides scaled dot-product attention with support for grouped query\n attention (GQA), causal masking, rotary position embeddings (RoPE), and\n dropout. *)\n\n(** {1:rope Rotary Position Embeddings} *)\n\nval rope : ?theta:float -> ?seq_dim:int -> (float, 'a) Nx.t -> (float, 'a) Nx.t\n(** [rope ?theta ?seq_dim x] applies rotary position embeddings to [x].\n\n [x] may have any rank [>= 2], with shape [[d0; ...; dn-1]] where:\n - [head_dim = dn-1] (last axis).\n - [seq_len] is on axis [seq_dim].\n\n [theta] defaults to [10000.0]. [seq_dim] defaults to [-2] (second-to-last\n axis). Negative [seq_dim] values are interpreted relative to rank.\n\n [head_dim] must be even.\n\n Raises [Invalid_argument] if [x] has rank < 2, if [seq_dim] is out of\n bounds, if [seq_dim] designates the last axis, or if [head_dim] is odd. *)\n\n(** {1:mha Multi-Head Attention} *)\n\nval attention_mask_key : string\n(** [attention_mask_key] is [\"attention_mask\"]. The well-known {!Context} key\n that {!multi_head_attention} reads during the forward pass. *)\n\nval multi_head_attention :\n embed_dim:int ->\n num_heads:int ->\n ?num_kv_heads:int ->\n ?dropout:float ->\n ?is_causal:bool ->\n ?rope:bool ->\n ?rope_theta:float ->\n unit ->\n (float, float) Layer.t\n(** [multi_head_attention ~embed_dim ~num_heads ()] is a multi-head\n self-attention layer.\n\n Input shape: [[batch; seq_len; embed_dim]]. Output shape:\n [[batch; seq_len; embed_dim]].\n\n [num_kv_heads] defaults to [num_heads] (standard MHA). When\n [num_kv_heads < num_heads], grouped query attention (GQA) is used.\n [num_heads] must be divisible by [num_kv_heads].\n\n [dropout] defaults to [0.0]. When positive, dropout is applied during\n training using keys from the implicit RNG scope.\n\n [is_causal] defaults to [false]. When [true], a causal mask prevents\n attending to future positions.\n\n [rope] defaults to [false]. When [true], rotary position embeddings are\n applied to Q and K before the attention computation. [rope_theta] defaults\n to [10000.0].\n\n When [ctx] contains {!attention_mask_key} (a bool or int32 tensor of shape\n [[batch; seq_k]]), it is applied as a padding mask. [true] / nonzero keeps\n the position, [false] / [0] masks it.\n\n Parameters:\n - [q_proj] ([[embed_dim; num_heads * head_dim]])\n - [k_proj] ([[embed_dim; num_kv_heads * head_dim]])\n - [v_proj] ([[embed_dim; num_kv_heads * head_dim]])\n - [out_proj] ([[num_heads * head_dim; embed_dim]])\n\n Raises [Invalid_argument] if [embed_dim] is not divisible by [num_heads], or\n if [num_heads] is not divisible by [num_kv_heads]. *)\n"
},
{
"path": "packages/kaun/lib/checkpoint.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\n\nlet shape_to_string s =\n \"[\" ^ String.concat \"; \" (Array.to_list (Array.map string_of_int s)) ^ \"]\"\n\nlet save path tree =\n let pairs = Ptree.flatten_with_paths tree in\n let items =\n List.map\n (fun (name, pt) ->\n let nx = Ptree.with_tensor pt { run = (fun t -> Nx_io.P t) } in\n (name, nx))\n pairs\n in\n Nx_io.save_safetensors path items\n\nlet load path ~like =\n let archive = Nx_io.load_safetensors path in\n let _, rebuild = Ptree.flatten like in\n let path_leaves = Ptree.flatten_with_paths like in\n let loaded =\n List.map\n (fun (name, template) ->\n match Hashtbl.find_opt archive name with\n | None -> invalid_argf \"Checkpoint.load: missing key %S\" name\n | Some (Nx_io.P nx) ->\n Ptree.with_tensor template\n {\n run =\n (fun tmpl ->\n let expected = Nx.shape tmpl in\n let actual = Nx.shape nx in\n if expected <> actual then\n invalid_argf\n \"Checkpoint.load: shape mismatch for %S: expected %s, \\\n got %s\"\n name (shape_to_string expected) (shape_to_string actual);\n let casted = Nx.cast (Nx.dtype tmpl) nx in\n Ptree.P casted);\n })\n path_leaves\n in\n rebuild loaded\n"
},
{
"path": "packages/kaun/lib/checkpoint.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Model checkpointing.\n\n {!Checkpoint} serializes {!Ptree.t} parameter trees to and from\n {{:https://huggingface.co/docs/safetensors/}SafeTensors} files. Tensor paths\n from {!Ptree.flatten_with_paths} become file keys (e.g.\n [\"layers.0.weight\"]). *)\n\nval save : string -> Ptree.t -> unit\n(** [save path t] writes [t]'s tensors to a safetensors file at [path].\n\n Raises [Failure] on I/O errors. *)\n\nval load : string -> like:Ptree.t -> Ptree.t\n(** [load path ~like] loads tensors from a safetensors file and reconstructs a\n tree with the same structure as [like].\n\n Each tensor is cast to [like]'s dtype if needed. Extra keys in the file are\n silently ignored.\n\n Raises [Invalid_argument] if a key required by [like] is missing from the\n file, or if a tensor's shape does not match [like]. Raises [Failure] on I/O\n errors. *)\n"
},
{
"path": "packages/kaun/lib/context.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Dtype = Nx_core.Dtype\n\ntype t = (string * Ptree.tensor) list\n\nlet empty = []\nlet set ~name tensor t = (name, tensor) :: t\nlet find t ~name = List.assoc_opt name t\n\nlet get_float_exn (type l) ~ctx t ~name ~(dtype : (float, l) Nx.dtype) :\n (float, l) Nx.t =\n match find t ~name with\n | Some (Ptree.P x) -> (\n match Dtype.equal_witness dtype (Nx.dtype x) with\n | Some Type.Equal -> x\n | None ->\n invalid_arg\n (Printf.sprintf \"%s: %s dtype mismatch (expected %s, got %s)\" ctx\n name (Dtype.to_string dtype)\n (Dtype.to_string (Nx.dtype x))))\n | None -> invalid_arg (Printf.sprintf \"%s: %s not found in context\" ctx name)\n\nlet get_int32_exn ~ctx t ~name : (int32, Bigarray.int32_elt) Nx.t =\n match find t ~name with\n | Some tensor -> Ptree.Tensor.to_typed_exn Nx.int32 tensor\n | None -> invalid_arg (Printf.sprintf \"%s: %s not found in context\" ctx name)\n\nlet get_bool_exn ~ctx t ~name : (bool, Nx.bool_elt) Nx.t =\n match find t ~name with\n | Some tensor -> Ptree.Tensor.to_typed_exn Nx.bool tensor\n | None -> invalid_arg (Printf.sprintf \"%s: %s not found in context\" ctx name)\n"
},
{
"path": "packages/kaun/lib/context.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Per-call auxiliary data for layers.\n\n A {!type:t} carries read-only tensors (attention masks, position ids,\n encoder memory) that specific layers consume during a forward pass. Most\n layers ignore the context; transformer layers read from it by well-known key\n names.\n\n {[\n let ctx =\n Context.empty\n |> Context.set ~name:\"attention_mask\" (Ptree.P mask)\n |> Context.set ~name:\"token_type_ids\" (Ptree.P ids)\n in\n Layer.apply model vars ~training:false ~ctx input_ids\n ]} *)\n\n(** {1:types Types} *)\n\ntype t\n(** The type for forward-pass contexts. *)\n\n(** {1:constructors Constructors} *)\n\nval empty : t\n(** [empty] is the empty context. *)\n\nval set : name:string -> Ptree.tensor -> t -> t\n(** [set ~name tensor ctx] is [ctx] with [name] bound to [tensor].\n\n Shadows any previous binding for [name]. *)\n\n(** {1:lookup Lookup} *)\n\nval find : t -> name:string -> Ptree.tensor option\n(** [find ctx ~name] is the tensor bound to [name] in [ctx], if any. *)\n\nval get_float_exn :\n ctx:string ->\n t ->\n name:string ->\n dtype:(float, 'l) Nx.dtype ->\n (float, 'l) Nx.t\n(** [get_float_exn ~ctx t ~name ~dtype] is the float tensor bound to [name],\n cast-checked against [dtype].\n\n Raises [Invalid_argument] if [name] is missing or has a different dtype.\n [ctx] is used in error messages. *)\n\nval get_int32_exn :\n ctx:string -> t -> name:string -> (int32, Bigarray.int32_elt) Nx.t\n(** [get_int32_exn ~ctx t ~name] is the int32 tensor bound to [name].\n\n Raises [Invalid_argument] if [name] is missing or has a different dtype. *)\n\nval get_bool_exn : ctx:string -> t -> name:string -> (bool, Nx.bool_elt) Nx.t\n(** [get_bool_exn ~ctx t ~name] is the bool tensor bound to [name].\n\n Raises [Invalid_argument] if [name] is missing or has a different dtype. *)\n"
},
{
"path": "packages/kaun/lib/data.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype 'a t = {\n next : unit -> 'a option;\n reset : unit -> unit;\n length : int option;\n}\n\n(* Constructors *)\n\nlet of_array a =\n let n = Array.length a in\n let i = ref 0 in\n {\n next =\n (fun () ->\n if !i >= n then None\n else\n let v = a.(!i) in\n incr i;\n Some v);\n reset = (fun () -> i := 0);\n length = Some n;\n }\n\nlet of_tensor t =\n let n = (Nx.shape t).(0) in\n let i = ref 0 in\n {\n next =\n (fun () ->\n if !i >= n then None\n else\n let v = Nx.slice [ I !i ] t in\n incr i;\n Some v);\n reset = (fun () -> i := 0);\n length = Some n;\n }\n\nlet of_tensors (x, y) =\n let nx = (Nx.shape x).(0) in\n let ny = (Nx.shape y).(0) in\n if nx <> ny then\n invalid_arg\n (Printf.sprintf \"Data.of_tensors: first dimensions differ (%d vs %d)\" nx\n ny);\n let n = nx in\n let i = ref 0 in\n {\n next =\n (fun () ->\n if !i >= n then None\n else\n let vx = Nx.slice [ I !i ] x in\n let vy = Nx.slice [ I !i ] y in\n incr i;\n Some (vx, vy));\n reset = (fun () -> i := 0);\n length = Some n;\n }\n\nlet of_fn n f =\n if n < 0 then\n invalid_arg (Printf.sprintf \"Data.of_fn: expected n >= 0, got %d\" n);\n let i = ref 0 in\n {\n next =\n (fun () ->\n if !i >= n then None\n else\n let v = f !i in\n incr i;\n Some v);\n reset = (fun () -> i := 0);\n length = Some n;\n }\n\nlet repeat n v = of_fn n (fun _ -> v)\n\n(* Transformers *)\n\nlet map f t =\n {\n next = (fun () -> Option.map f (t.next ()));\n reset = t.reset;\n length = t.length;\n }\n\nlet batch ?(drop_last = false) n t =\n if n <= 0 then\n invalid_arg (Printf.sprintf \"Data.batch: expected n > 0, got %d\" n);\n let batch_len =\n Option.map (fun l -> if drop_last then l / n else (l + n - 1) / n) t.length\n in\n {\n next =\n (fun () ->\n match t.next () with\n | None -> None\n | Some first ->\n let buf = Array.make n first in\n let k = ref 1 in\n let continue = ref true in\n while !k < n && !continue do\n match t.next () with\n | Some v ->\n buf.(!k) <- v;\n incr k\n | None -> continue := false\n done;\n if !k < n && drop_last then None\n else if !k < n then Some (Array.sub buf 0 !k)\n else Some buf);\n reset = t.reset;\n length = batch_len;\n }\n\nlet map_batch ?drop_last n f t = map f (batch ?drop_last n t)\n\nlet shuffle t =\n match t.length with\n | None -> invalid_arg \"Data.shuffle: requires a pipeline with known length\"\n | Some n ->\n let perm_tensor = Nx.permutation n in\n let perm = Array.map Int32.to_int (Nx.to_array perm_tensor) in\n (* Eagerly materialize the upstream into an array *)\n let elements =\n Array.init n (fun _ ->\n match t.next () with Some v -> v | None -> assert false)\n in\n let i = ref 0 in\n {\n next =\n (fun () ->\n if !i >= n then None\n else\n let v = elements.(perm.(!i)) in\n incr i;\n Some v);\n reset = (fun () -> i := 0);\n length = Some n;\n }\n\n(* Consumers *)\n\nlet iter f t =\n let rec loop () =\n match t.next () with\n | None -> ()\n | Some v ->\n f v;\n loop ()\n in\n loop ()\n\nlet iteri f t =\n let i = ref 0 in\n let rec loop () =\n match t.next () with\n | None -> ()\n | Some v ->\n f !i v;\n incr i;\n loop ()\n in\n loop ()\n\nlet fold f init t =\n let rec loop acc =\n match t.next () with None -> acc | Some v -> loop (f acc v)\n in\n loop init\n\nlet to_array t =\n let items = ref [] in\n iter (fun v -> items := v :: !items) t;\n Array.of_list (List.rev !items)\n\nlet rec to_seq t () =\n match t.next () with None -> Seq.Nil | Some v -> Seq.Cons (v, to_seq t)\n\n(* Properties *)\n\nlet reset t = t.reset ()\nlet length t = t.length\n\n(* Utilities *)\n\nlet stack_batch tensors = Nx.stack (Array.to_list tensors)\nlet shuffle_pipeline = shuffle\n\nlet prepare ?(shuffle = false) ~batch_size ?(drop_last = true) (x, y) =\n let nx = (Nx.shape x).(0) in\n let ny = (Nx.shape y).(0) in\n if nx <> ny then\n invalid_arg\n (Printf.sprintf \"Data.prepare: first dimensions differ (%d vs %d)\" nx ny);\n if batch_size <= 0 then\n invalid_arg\n (Printf.sprintf \"Data.prepare: expected batch_size > 0, got %d\" batch_size);\n let indices = of_fn nx Fun.id in\n let indices = if shuffle then shuffle_pipeline indices else indices in\n map_batch ~drop_last batch_size\n (fun idx_arr ->\n let n = Array.length idx_arr in\n let xs = Array.init n (fun j -> Nx.slice [ I idx_arr.(j) ] x) in\n let ys = Array.init n (fun j -> Nx.slice [ I idx_arr.(j) ] y) in\n (stack_batch xs, stack_batch ys))\n indices\n"
},
{
"path": "packages/kaun/lib/data.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Lazy, composable data pipelines for training.\n\n A {!type:t} is a resettable iterator over elements of type ['a]. Pipelines\n are built by composing constructors, transformers, and consumers.\n\n {[\n Data.of_array examples |> Data.shuffle |> Data.map_batch 32 collate\n |> Data.iter train_step\n ]} *)\n\n(** {1:types Types} *)\n\ntype 'a t\n(** The type for lazy data pipelines producing elements of type ['a]. *)\n\n(** {1:constructors Constructors} *)\n\nval of_array : 'a array -> 'a t\n(** [of_array a] is a pipeline yielding the elements of [a] in order. *)\n\nval of_tensor : ('a, 'b) Nx.t -> ('a, 'b) Nx.t t\n(** [of_tensor t] is a pipeline yielding slices along the first dimension of\n [t]. Each element has shape [t.shape[1:]]. *)\n\nval of_tensors :\n ('a, 'b) Nx.t * ('c, 'd) Nx.t -> (('a, 'b) Nx.t * ('c, 'd) Nx.t) t\n(** [of_tensors (x, y)] is a pipeline yielding paired slices along the first\n dimension of [x] and [y].\n\n Raises [Invalid_argument] if [x] and [y] have different first dimension\n sizes. *)\n\nval of_fn : int -> (int -> 'a) -> 'a t\n(** [of_fn n f] is a pipeline yielding [f 0], [f 1], ..., [f (n - 1)].\n\n Raises [Invalid_argument] if [n < 0]. *)\n\nval repeat : int -> 'a -> 'a t\n(** [repeat n v] is a pipeline that yields [v] exactly [n] times.\n\n Raises [Invalid_argument] if [n < 0]. *)\n\n(** {1:transformers Transformers} *)\n\nval map : ('a -> 'b) -> 'a t -> 'b t\n(** [map f t] is a pipeline that applies [f] to each element of [t]. *)\n\nval batch : ?drop_last:bool -> int -> 'a t -> 'a array t\n(** [batch ?drop_last n t] is a pipeline yielding arrays of [n] consecutive\n elements from [t].\n\n [drop_last] defaults to [false]. When [true], the final batch is dropped if\n it has fewer than [n] elements.\n\n Raises [Invalid_argument] if [n <= 0]. *)\n\nval map_batch : ?drop_last:bool -> int -> ('a array -> 'b) -> 'a t -> 'b t\n(** [map_batch ?drop_last n f t] is [map f (batch ?drop_last n t)]. *)\n\nval shuffle : 'a t -> 'a t\n(** [shuffle t] is a pipeline that yields the elements of [t] in a random order.\n The permutation is computed once when the pipeline is created.\n\n Random keys are drawn from the implicit RNG scope.\n\n Raises [Invalid_argument] if [t] has unknown length. *)\n\n(** {1:consumers Consumers} *)\n\nval iter : ('a -> unit) -> 'a t -> unit\n(** [iter f t] applies [f] to each element of [t]. *)\n\nval iteri : (int -> 'a -> unit) -> 'a t -> unit\n(** [iteri f t] applies [f i x] to each element [x] of [t], where [i] is the\n 0-based index. *)\n\nval fold : ('acc -> 'a -> 'acc) -> 'acc -> 'a t -> 'acc\n(** [fold f init t] folds [f] over the elements of [t]. *)\n\nval to_array : 'a t -> 'a array\n(** [to_array t] collects all elements of [t] into an array. *)\n\nval to_seq : 'a t -> 'a Seq.t\n(** [to_seq t] is a standard [Seq.t] view of [t]. Does not reset [t]. *)\n\n(** {1:properties Properties} *)\n\nval reset : 'a t -> unit\n(** [reset t] resets [t] so that iteration starts from the beginning. *)\n\nval length : 'a t -> int option\n(** [length t] is the number of elements in [t], if known. *)\n\n(** {1:utilities Utilities} *)\n\nval stack_batch : ('a, 'b) Nx.t array -> ('a, 'b) Nx.t\n(** [stack_batch tensors] stacks an array of tensors along a new first axis.\n Equivalent to [Nx.stack (Array.to_list tensors)]. *)\n\nval prepare :\n ?shuffle:bool ->\n batch_size:int ->\n ?drop_last:bool ->\n ('a, 'b) Nx.t * ('c, 'd) Nx.t ->\n (('a, 'b) Nx.t * ('c, 'd) Nx.t) t\n(** [prepare ?shuffle ~batch_size (x, y)] is a pipeline that yields batched\n tensor pairs from [x] and [y].\n\n Each yielded pair has shape [[batch_size; ...]] along the first dimension.\n\n [shuffle] defaults to [false]. When [true], elements are yielded in a random\n order. [drop_last] defaults to [true].\n\n Raises [Invalid_argument] if [x] and [y] have different first dimension\n sizes, or if [batch_size <= 0]. *)\n"
},
{
"path": "packages/kaun/lib/datasets/cifar10.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Bigarray\nopen Dataset_utils\n\nlet src = Logs.Src.create \"kaun.datasets.cifar10\" ~doc:\"CIFAR-10 dataset loader\"\n\nmodule Log = (val Logs.src_log src : Logs.LOG)\n\nmodule Config = struct\n let url = \"https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz\"\n let cache_subdir = \"cifar10/\"\n let archive_name = \"cifar-10-binary.tar.gz\"\n let extracted_subdir = \"cifar-10-batches-bin/\"\n let height = 32\n let width = 32\n let channels = 3\n let image_size = channels * height * width\n let entry_size = 1 + image_size\n let entries_per_batch = 10000\n\n let train_batches =\n [\n \"data_batch_1.bin\";\n \"data_batch_2.bin\";\n \"data_batch_3.bin\";\n \"data_batch_4.bin\";\n \"data_batch_5.bin\";\n ]\n\n let test_batches = [ \"test_batch.bin\" ]\nend\n\nlet ensure_dataset () =\n let dataset_dir = get_cache_dir Config.cache_subdir in\n mkdir_p dataset_dir;\n let archive_path = dataset_dir ^ Config.archive_name in\n let extracted_dir = dataset_dir ^ Config.extracted_subdir in\n let check_file = extracted_dir ^ \"test_batch.bin\" in\n if not (Sys.file_exists check_file) then (\n ensure_file Config.url archive_path;\n if\n not\n (ensure_extracted_tar_gz ~tar_gz_path:archive_path\n ~target_dir:dataset_dir ~check_file)\n then\n failwith\n (Printf.sprintf \"Failed to extract CIFAR-10 archive to %s\" extracted_dir));\n extracted_dir\n\nlet read_batch_file ~extracted_dir filename =\n let path = extracted_dir ^ filename in\n Log.debug (fun m -> m \"Reading CIFAR-10 batch: %s\" path);\n let ic = open_in_bin path in\n Fun.protect\n ~finally:(fun () -> close_in ic)\n (fun () ->\n let s = really_input_string ic (in_channel_length ic) in\n let num_entries = String.length s / Config.entry_size in\n if String.length s <> num_entries * Config.entry_size then\n failwith\n (Printf.sprintf\n \"CIFAR-10 batch %s has unexpected size %d (expected multiple of \\\n %d)\"\n filename (String.length s) Config.entry_size);\n (s, num_entries))\n\nlet load () =\n let extracted_dir = ensure_dataset () in\n let load_split batch_files expected_total =\n let images =\n Genarray.create int8_unsigned c_layout\n [| expected_total; Config.channels; Config.height; Config.width |]\n in\n let labels = Array1.create int8_unsigned c_layout expected_total in\n let flat = Bigarray.reshape_1 images (expected_total * Config.image_size) in\n let offset = ref 0 in\n List.iter\n (fun filename ->\n let s, num_entries = read_batch_file ~extracted_dir filename in\n for i = 0 to num_entries - 1 do\n let entry_offset = i * Config.entry_size in\n let idx = !offset + i in\n Array1.unsafe_set labels idx (Char.code s.[entry_offset]);\n let img_offset = entry_offset + 1 in\n let base = idx * Config.image_size in\n for p = 0 to Config.image_size - 1 do\n Array1.unsafe_set flat (base + p)\n (Char.code (String.unsafe_get s (img_offset + p)))\n done\n done;\n offset := !offset + num_entries)\n batch_files;\n (images, labels)\n in\n Log.info (fun m -> m \"Loading CIFAR-10 datasets...\");\n let train_images, train_labels =\n load_split Config.train_batches\n (List.length Config.train_batches * Config.entries_per_batch)\n in\n let test_images, test_labels =\n load_split Config.test_batches\n (List.length Config.test_batches * Config.entries_per_batch)\n in\n Log.info (fun m -> m \"CIFAR-10 loading complete\");\n ((train_images, train_labels), (test_images, test_labels))\n"
},
{
"path": "packages/kaun/lib/datasets/dataset_utils.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet src = Logs.Src.create \"kaun.datasets\" ~doc:\"Kaun datasets\"\n\nmodule Log = (val Logs.src_log src : Logs.LOG)\n\nlet mkdir_p path =\n if path = \"\" || path = \".\" || path = Filename.dir_sep then ()\n else\n let components =\n String.split_on_char Filename.dir_sep.[0] path |> List.filter (( <> ) \"\")\n in\n let is_absolute = path <> \"\" && path.[0] = Filename.dir_sep.[0] in\n let initial_prefix = if is_absolute then Filename.dir_sep else \".\" in\n ignore\n (List.fold_left\n (fun prefix comp ->\n let next =\n if prefix = Filename.dir_sep then Filename.dir_sep ^ comp\n else Filename.concat prefix comp\n in\n (if Sys.file_exists next then (\n if not (Sys.is_directory next) then\n failwith\n (Printf.sprintf \"mkdir_p: '%s' exists but is not a directory\"\n next))\n else\n try Unix.mkdir next 0o755\n with Unix.Unix_error (Unix.EEXIST, _, _) ->\n if not (Sys.is_directory next) then\n failwith\n (Printf.sprintf\n \"mkdir_p: '%s' appeared as non-directory after EEXIST\"\n next));\n next)\n initial_prefix components)\n\nlet get_cache_dir ?(getenv = Sys.getenv_opt) dataset_name =\n let root =\n match getenv \"RAVEN_CACHE_ROOT\" with\n | Some dir when dir <> \"\" -> dir\n | _ ->\n let xdg =\n match getenv \"XDG_CACHE_HOME\" with\n | Some d when d <> \"\" -> d\n | _ -> Filename.concat (Sys.getenv \"HOME\") \".cache\"\n in\n Filename.concat xdg \"raven\"\n in\n let path =\n List.fold_left Filename.concat root (\"datasets\" :: [ dataset_name ])\n in\n let sep = Filename.dir_sep.[0] in\n if path <> \"\" && path.[String.length path - 1] = sep then path\n else path ^ Filename.dir_sep\n\nlet curl_download ~url ~dest () =\n let check =\n lazy (Unix.system \"command -v curl >/dev/null 2>&1\" = Unix.WEXITED 0)\n in\n if not (Lazy.force check) then failwith \"curl not found on PATH\";\n mkdir_p (Filename.dirname dest);\n let cmd =\n Printf.sprintf \"curl -L --fail -s -o %s %s\" (Filename.quote dest)\n (Filename.quote url)\n in\n match Unix.system cmd with\n | Unix.WEXITED 0 -> ()\n | _ ->\n (try Sys.remove dest with Sys_error _ -> ());\n failwith (Printf.sprintf \"Failed to download %s\" url)\n\nlet download_file url dest_path =\n Log.info (fun m -> m \"Downloading %s to %s\" (Filename.basename url) dest_path);\n curl_download ~url ~dest:dest_path ();\n Log.info (fun m -> m \"Downloaded %s\" (Filename.basename dest_path))\n\nlet ensure_file url dest_path =\n if not (Sys.file_exists dest_path) then download_file url dest_path\n else Log.debug (fun m -> m \"Found %s\" dest_path)\n\nlet ensure_decompressed_gz ~gz_path ~target_path =\n if Sys.file_exists target_path then (\n Log.debug (fun m -> m \"Found %s\" target_path);\n true)\n else if Sys.file_exists gz_path then (\n Log.info (fun m -> m \"Decompressing %s...\" gz_path);\n let ic = Gzip.open_in gz_path in\n let oc = open_out_bin target_path in\n Fun.protect\n ~finally:(fun () ->\n Gzip.close_in ic;\n close_out oc)\n (fun () ->\n let buf = Bytes.create 4096 in\n let rec loop () =\n let n = Gzip.input ic buf 0 4096 in\n if n > 0 then (\n output oc buf 0 n;\n loop ())\n in\n loop ());\n Log.info (fun m -> m \"Decompressed to %s\" target_path);\n true)\n else (\n Log.warn (fun m -> m \"Compressed file %s not found\" gz_path);\n false)\n\nlet ensure_extracted_tar_gz ~tar_gz_path ~target_dir ~check_file =\n if Sys.file_exists check_file then (\n Log.debug (fun m -> m \"Found %s\" check_file);\n true)\n else if Sys.file_exists tar_gz_path then (\n Log.info (fun m -> m \"Extracting %s...\" tar_gz_path);\n mkdir_p target_dir;\n let cmd =\n Printf.sprintf \"tar -xzf %s -C %s\"\n (Filename.quote tar_gz_path)\n (Filename.quote target_dir)\n in\n match Unix.system cmd with\n | Unix.WEXITED 0 ->\n Log.info (fun m -> m \"Extracted to %s\" target_dir);\n true\n | _ ->\n Log.warn (fun m -> m \"Failed to extract %s\" tar_gz_path);\n false)\n else (\n Log.warn (fun m -> m \"Archive %s not found\" tar_gz_path);\n false)\n"
},
{
"path": "packages/kaun/lib/datasets/dune",
"content": "(library\n (name kaun_datasets)\n (public_name kaun.datasets)\n (libraries unix zip rune nx kaun logs))\n"
},
{
"path": "packages/kaun/lib/datasets/kaun_datasets.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet mnist ?(fashion = false) ?(normalize = true) ?(data_format = `NCHW) () =\n let (train_images, train_labels), (test_images, test_labels) =\n Mnist.load ~fashion_mnist:fashion\n in\n let make_tensors images labels =\n let n = Bigarray.Array3.dim1 images in\n let h = Bigarray.Array3.dim2 images in\n let w = Bigarray.Array3.dim3 images in\n let x =\n Nx.of_bigarray (Bigarray.genarray_of_array3 images)\n |> Nx.reshape [| n; h; w; 1 |]\n |> Nx.cast Nx.float32\n in\n let x = if normalize then Nx.div_s x 255.0 else x in\n let x =\n match data_format with\n | `NCHW -> Nx.transpose x ~axes:[ 0; 3; 1; 2 ]\n | `NHWC -> x\n in\n let y =\n Nx.of_bigarray (Bigarray.genarray_of_array1 labels) |> Nx.cast Nx.int32\n in\n (x, y)\n in\n let train = make_tensors train_images train_labels in\n let test = make_tensors test_images test_labels in\n (train, test)\n\nlet cifar10 ?(normalize = true) ?(data_format = `NCHW) () =\n let (train_images, train_labels), (test_images, test_labels) =\n Cifar10.load ()\n in\n let make_tensors images labels =\n let x = Nx.of_bigarray images |> Nx.cast Nx.float32 in\n let x = if normalize then Nx.div_s x 255.0 else x in\n let x =\n match data_format with\n | `NCHW -> x\n | `NHWC -> Nx.transpose x ~axes:[ 0; 2; 3; 1 ]\n in\n let y =\n Nx.of_bigarray (Bigarray.genarray_of_array1 labels) |> Nx.cast Nx.int32\n in\n (x, y)\n in\n let train = make_tensors train_images train_labels in\n let test = make_tensors test_images test_labels in\n (train, test)\n"
},
{
"path": "packages/kaun/lib/datasets/kaun_datasets.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Dataset loaders for kaun.\n\n Datasets are downloaded on demand and cached locally under\n [$RAVEN_CACHE_ROOT/datasets/] (or [$XDG_CACHE_HOME/raven/datasets/]). *)\n\nval mnist :\n ?fashion:bool ->\n ?normalize:bool ->\n ?data_format:[ `NCHW | `NHWC ] ->\n unit ->\n (Nx.float32_t * Nx.int32_t) * (Nx.float32_t * Nx.int32_t)\n(** [mnist ()] is [((x_train, y_train), (x_test, y_test))].\n\n Images are float32 in \\[0, 1\\] (when [normalize] is [true], the default).\n Labels are int32 class indices.\n\n [fashion] selects Fashion-MNIST when [true]. Defaults to [false].\n [data_format] defaults to [`NCHW].\n\n Tensor shapes:\n - [`NCHW]: images [[N; 1; 28; 28]], labels [[N]]\n - [`NHWC]: images [[N; 28; 28; 1]], labels [[N]]\n\n Raises [Failure] on download or parsing errors. *)\n\nval cifar10 :\n ?normalize:bool ->\n ?data_format:[ `NCHW | `NHWC ] ->\n unit ->\n (Nx.float32_t * Nx.int32_t) * (Nx.float32_t * Nx.int32_t)\n(** [cifar10 ()] is [((x_train, y_train), (x_test, y_test))].\n\n Images are float32 in \\[0, 1\\] (when [normalize] is [true], the default).\n Labels are int32 class indices (0--9: airplane, automobile, bird, cat, deer,\n dog, frog, horse, ship, truck).\n\n [data_format] defaults to [`NCHW].\n\n Tensor shapes:\n - [`NCHW]: images [[N; 3; 32; 32]], labels [[N]]\n - [`NHWC]: images [[N; 32; 32; 3]], labels [[N]]\n\n Raises [Failure] on download, extraction, or parsing errors. *)\n"
},
{
"path": "packages/kaun/lib/datasets/mnist.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Bigarray\nopen Dataset_utils\n\nlet src = Logs.Src.create \"kaun.datasets.mnist\" ~doc:\"MNIST dataset loader\"\n\nmodule Log = (val Logs.src_log src : Logs.LOG)\n\nmodule Config = struct\n type t = {\n name : string;\n cache_subdir : string;\n train_images_url : string;\n train_labels_url : string;\n test_images_url : string;\n test_labels_url : string;\n image_magic_number : int;\n label_magic_number : int;\n }\n\n let mnist =\n {\n name = \"MNIST\";\n cache_subdir = \"mnist/\";\n train_images_url =\n \"https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz\";\n train_labels_url =\n \"https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz\";\n test_images_url =\n \"https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz\";\n test_labels_url =\n \"https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz\";\n image_magic_number = 2051;\n label_magic_number = 2049;\n }\n\n let fashion_mnist =\n {\n name = \"Fashion-MNIST\";\n cache_subdir = \"fashion-mnist/\";\n train_images_url =\n \"http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz\";\n train_labels_url =\n \"http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz\";\n test_images_url =\n \"http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz\";\n test_labels_url =\n \"http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz\";\n image_magic_number = 2051;\n label_magic_number = 2049;\n }\nend\n\nlet read_int32_be s pos =\n let b1 = Char.code s.[pos] in\n let b2 = Char.code s.[pos + 1] in\n let b3 = Char.code s.[pos + 2] in\n let b4 = Char.code s.[pos + 3] in\n (b1 lsl 24) lor (b2 lsl 16) lor (b3 lsl 8) lor b4\n\nlet ensure_dataset config =\n let dataset_dir = get_cache_dir config.Config.cache_subdir in\n mkdir_p dataset_dir;\n let files_to_process =\n [\n (\"train-images-idx3-ubyte\", config.Config.train_images_url);\n (\"train-labels-idx1-ubyte\", config.Config.train_labels_url);\n (\"t10k-images-idx3-ubyte\", config.Config.test_images_url);\n (\"t10k-labels-idx1-ubyte\", config.Config.test_labels_url);\n ]\n in\n List.iter\n (fun (base_filename, url) ->\n let gz_filename = base_filename ^ \".gz\" in\n let gz_path = dataset_dir ^ gz_filename in\n let path = dataset_dir ^ base_filename in\n if not (Sys.file_exists path) then (\n Log.debug (fun m ->\n m \"File %s not found for %s dataset\" base_filename config.name);\n ensure_file url gz_path;\n if not (ensure_decompressed_gz ~gz_path ~target_path:path) then\n failwith (Printf.sprintf \"Failed to obtain decompressed file %s\" path))\n else Log.debug (fun m -> m \"Found decompressed file %s\" path))\n files_to_process\n\nlet read_idx_file ~read_header ~create_array ~populate_array ~expected_magic\n config filename =\n Log.debug (fun m -> m \"Reading %s file: %s\" config.Config.name filename);\n let ic = open_in_bin filename in\n let s =\n Fun.protect\n ~finally:(fun () -> close_in ic)\n (fun () -> really_input_string ic (in_channel_length ic))\n in\n let magic = read_int32_be s 0 in\n if magic <> expected_magic then\n failwith\n (Printf.sprintf \"Invalid magic number %d in %s (expected %d)\" magic\n filename expected_magic);\n let dimensions, data_offset = read_header s in\n let total_items, data_len =\n match dimensions with\n | [| d1 |] -> (d1, d1)\n | [| d1; d2; d3 |] -> (d1, d1 * d2 * d3)\n | _ -> failwith \"Unsupported dimension format\"\n in\n let expected_len = data_offset + data_len in\n if String.length s <> expected_len then\n failwith\n (Printf.sprintf\n \"File %s has unexpected length: %d vs %d (header offset %d, data len \\\n %d)\"\n filename (String.length s) expected_len data_offset data_len);\n let arr = create_array dimensions in\n populate_array arr s data_offset total_items;\n arr\n\nlet read_images config filename =\n let read_header s =\n let num_images = read_int32_be s 4 in\n let num_rows = read_int32_be s 8 in\n let num_cols = read_int32_be s 12 in\n ([| num_images; num_rows; num_cols |], 16)\n in\n let create_array dims =\n Array3.create int8_unsigned c_layout dims.(0) dims.(1) dims.(2)\n in\n let populate_array arr s offset _total_items =\n let num_images = Array3.dim1 arr in\n let num_rows = Array3.dim2 arr in\n let num_cols = Array3.dim3 arr in\n let img_size = num_rows * num_cols in\n for i = 0 to num_images - 1 do\n let start_pos = offset + (i * img_size) in\n for r = 0 to num_rows - 1 do\n for c = 0 to num_cols - 1 do\n let pos = start_pos + (r * num_cols) + c in\n arr.{i, r, c} <- Char.code s.[pos]\n done\n done\n done\n in\n read_idx_file ~read_header ~create_array ~populate_array\n ~expected_magic:config.Config.image_magic_number config filename\n\nlet read_labels config filename =\n let read_header s =\n let num_labels = read_int32_be s 4 in\n ([| num_labels |], 8)\n in\n let create_array dims = Array1.create int8_unsigned c_layout dims.(0) in\n let populate_array arr s offset total_items =\n for i = 0 to total_items - 1 do\n arr.{i} <- Char.code s.[offset + i]\n done\n in\n read_idx_file ~read_header ~create_array ~populate_array\n ~expected_magic:config.Config.label_magic_number config filename\n\nlet load ~fashion_mnist =\n let config = if fashion_mnist then Config.fashion_mnist else Config.mnist in\n ensure_dataset config;\n let dataset_dir = get_cache_dir config.Config.cache_subdir in\n let train_images_path = dataset_dir ^ \"train-images-idx3-ubyte\" in\n let train_labels_path = dataset_dir ^ \"train-labels-idx1-ubyte\" in\n let test_images_path = dataset_dir ^ \"t10k-images-idx3-ubyte\" in\n let test_labels_path = dataset_dir ^ \"t10k-labels-idx1-ubyte\" in\n Log.info (fun m -> m \"Loading %s datasets...\" config.name);\n let train_images = read_images config train_images_path in\n let train_labels = read_labels config train_labels_path in\n let test_images = read_images config test_images_path in\n let test_labels = read_labels config test_labels_path in\n Log.info (fun m -> m \"%s loading complete\" config.name);\n ((train_images, train_labels), (test_images, test_labels))\n"
},
{
"path": "packages/kaun/lib/dune",
"content": "(library\n (name kaun)\n (public_name kaun)\n (libraries rune vega nx nx.core nx.io))\n"
},
{
"path": "packages/kaun/lib/fn.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\n\nlet invalid_argf_fn fn fmt =\n Printf.ksprintf (fun msg -> invalid_argf \"Fn.%s: %s\" fn msg) fmt\n\n(* Helpers *)\n\nlet normalize_axis ~fn ~ndim ax =\n let axis = if ax < 0 then ndim + ax else ax in\n if axis < 0 || axis >= ndim then\n invalid_argf_fn fn \"axis %d out of bounds for rank %d\" ax ndim;\n axis\n\nlet normalize_axes ~fn ~ndim axes =\n match axes with\n | [] -> invalid_argf_fn fn \"axes must contain at least one axis\"\n | lst -> List.map (normalize_axis ~fn ~ndim) lst\n\nlet keep_shape ~axes x_shape =\n Array.mapi\n (fun idx dim -> if List.exists (fun ax -> ax = idx) axes then 1 else dim)\n x_shape\n\nlet unaffected_axes ~ndim ~axes =\n Array.init ndim Fun.id |> Array.to_list\n |> List.filter (fun ax -> not (List.exists (( = ) ax) axes))\n\nlet core_shape ~axes:unaffected x_shape =\n Array.of_list (List.map (fun ax -> x_shape.(ax)) unaffected)\n\nlet broadcast_param ~fn ~name ~x_shape ~keep_shape ~core_shape param =\n let param_shape = Nx.shape param in\n if param_shape = keep_shape then param\n else if param_shape = core_shape then Nx.reshape keep_shape param\n else if param_shape = x_shape then param\n else\n invalid_argf_fn fn \"%s: shape must match normalized axes or remaining axes\"\n name\n\n(* Normalization *)\n\nlet batch_norm ?axes ?(epsilon = 1e-5) ~scale ~bias x =\n let ndim = Nx.ndim x in\n let axes =\n let default =\n match axes with\n | Some ax -> ax\n | None ->\n if ndim = 2 then [ 0 ] else if ndim = 4 then [ 0; 2; 3 ] else [ 0 ]\n in\n normalize_axes ~fn:\"batch_norm\" ~ndim default\n in\n let x_shape = Nx.shape x in\n let keep = keep_shape ~axes x_shape in\n let unaffected = unaffected_axes ~ndim ~axes in\n let core = core_shape ~axes:unaffected x_shape in\n let broadcast name param =\n let param =\n if Nx.dtype param <> Nx.dtype x then Nx.cast (Nx.dtype x) param else param\n in\n broadcast_param ~fn:\"batch_norm\" ~name ~x_shape ~keep_shape:keep\n ~core_shape:core param\n in\n let mean_x = Nx.mean x ~axes ~keepdims:true in\n let variance = Nx.var x ~axes ~keepdims:true in\n let eps = Nx.scalar_like x epsilon in\n let normalized = Nx.mul (Nx.sub x mean_x) (Nx.rsqrt (Nx.add variance eps)) in\n let scale_b = broadcast \"scale\" scale in\n let bias_b = broadcast \"bias\" bias in\n Nx.add (Nx.mul normalized scale_b) bias_b\n\nlet rms_norm ?axes ?(epsilon = 1e-5) ?gamma x =\n let ndim = Nx.ndim x in\n let axes =\n let default = match axes with Some ax -> ax | None -> [ -1 ] in\n normalize_axes ~fn:\"rms_norm\" ~ndim default\n in\n let x_shape = Nx.shape x in\n let keep = keep_shape ~axes x_shape in\n let mean_square = Nx.mean (Nx.mul x x) ~axes ~keepdims:true in\n let eps = Nx.scalar_like x epsilon in\n let normalized = Nx.mul x (Nx.rsqrt (Nx.add mean_square eps)) in\n match gamma with\n | None -> normalized\n | Some gamma ->\n let gamma_shape = Nx.shape gamma in\n let gamma =\n if gamma_shape = keep then gamma\n else\n let unaffected = unaffected_axes ~ndim ~axes in\n let core = core_shape ~axes:unaffected x_shape in\n if gamma_shape = core then Nx.reshape keep gamma\n else if gamma_shape = x_shape then gamma\n else\n invalid_argf_fn \"rms_norm\"\n \"gamma: shape must match normalized axes or remaining axes\"\n in\n Nx.mul normalized gamma\n\nlet layer_norm ?(axes = [ -1 ]) ?(epsilon = 1e-5) ?gamma ?beta x =\n let ndim = Nx.ndim x in\n let axes =\n List.map\n (fun ax ->\n let axis = if ax < 0 then ndim + ax else ax in\n if axis < 0 || axis >= ndim then\n invalid_argf_fn \"layer_norm\" \"axis %d out of bounds for rank %d\" ax\n ndim;\n axis)\n axes\n in\n let x_shape = Nx.shape x in\n let keep =\n Array.mapi (fun idx dim -> if List.mem idx axes then dim else 1) x_shape\n in\n let broadcast_param name param =\n let param_shape = Nx.shape param in\n if param_shape = x_shape then param\n else if param_shape = keep then param\n else\n let axes_shape = Array.of_list (List.map (fun ax -> x_shape.(ax)) axes) in\n if param_shape = axes_shape then Nx.reshape keep param\n else\n invalid_argf_fn \"layer_norm\" \"%s: shape must match normalized axes\" name\n in\n let mean_x = Nx.mean x ~axes ~keepdims:true in\n let centered = Nx.sub x mean_x in\n let variance = Nx.mean (Nx.mul centered centered) ~axes ~keepdims:true in\n let eps = Nx.scalar_like x epsilon in\n let inv_std = Nx.rsqrt (Nx.add variance eps) in\n let normalized = Nx.mul centered inv_std in\n let with_scale =\n match gamma with\n | None -> normalized\n | Some gamma ->\n let gamma_broadcast = broadcast_param \"gamma\" gamma in\n Nx.mul normalized gamma_broadcast\n in\n match beta with\n | None -> with_scale\n | Some beta ->\n let beta_broadcast = broadcast_param \"beta\" beta in\n Nx.add with_scale beta_broadcast\n\n(* Embedding *)\n\nlet embedding ?(scale = true) ~embedding indices =\n let embed_shape = Nx.shape embedding in\n if Array.length embed_shape <> 2 then\n invalid_argf_fn \"embedding\"\n \"embedding matrix must have shape [vocab_size; embed_dim]\";\n let embed_dim = embed_shape.(1) in\n let indices_shape = Nx.shape indices in\n let is_scalar = Array.length indices_shape = 0 in\n let vocab_size = embed_shape.(0) in\n if vocab_size <= 0 then\n invalid_argf_fn \"embedding\" \"vocabulary dimension must be positive\";\n let flat_size = Array.fold_left ( * ) 1 indices_shape in\n let indices_flat =\n if is_scalar then Nx.reshape [| 1 |] indices\n else Nx.reshape [| flat_size |] indices\n in\n let gathered = Nx.take ~axis:0 indices_flat embedding in\n let output_shape =\n if is_scalar then [| embed_dim |]\n else Array.append indices_shape [| embed_dim |]\n in\n let embedded = Nx.reshape output_shape gathered in\n if not scale then embedded\n else\n let factor =\n Nx.scalar_like embedding (Stdlib.sqrt (float_of_int embed_dim))\n in\n Nx.mul embedded factor\n\n(* Dropout *)\n\nlet dropout ~rate x =\n if rate < 0.0 || rate >= 1.0 then\n invalid_argf_fn \"dropout\" \"rate must satisfy 0.0 <= rate < 1.0\";\n let tensor_dtype = Nx.dtype x in\n if not (Nx_core.Dtype.is_float tensor_dtype) then\n invalid_argf_fn \"dropout\" \"requires floating point dtype\";\n if rate = 0.0 then x\n else\n let keep_prob = 1.0 -. rate in\n let random_vals = Nx.rand tensor_dtype (Nx.shape x) in\n let threshold = Nx.scalar_like x keep_prob in\n let keep_mask = Nx.less random_vals threshold in\n let keep_mask_float = Nx.cast tensor_dtype keep_mask in\n let scale = Nx.scalar_like x (1.0 /. keep_prob) in\n Nx.mul x (Nx.mul keep_mask_float scale)\n\n(* Attention *)\n\nlet dot_product_attention (type b) ?attention_mask ?scale ?dropout_rate\n ?(is_causal = false) (q : (float, b) Nx.t) (k : (float, b) Nx.t)\n (v : (float, b) Nx.t) =\n let check_float name (t : (float, b) Nx.t) =\n match Nx.dtype t with\n | Nx.Float16 -> ()\n | Nx.Float32 -> ()\n | Nx.Float64 -> ()\n | _ ->\n invalid_argf_fn \"dot_product_attention\"\n \"%s: requires floating point dtype\" name\n in\n check_float \"query\" q;\n check_float \"key\" k;\n check_float \"value\" v;\n let q_shape = Nx.shape q in\n let k_shape = Nx.shape k in\n let v_shape = Nx.shape v in\n let q_rank = Array.length q_shape in\n if q_rank < 2 then\n invalid_argf_fn \"dot_product_attention\" \"query: must have rank >= 2\";\n if Array.length k_shape <> q_rank || Array.length v_shape <> q_rank then\n invalid_argf_fn \"dot_product_attention\" \"key/value: must match query rank\";\n let depth = q_shape.(q_rank - 1) in\n if k_shape.(q_rank - 1) <> depth then\n invalid_argf_fn \"dot_product_attention\"\n \"key last dim %d does not match query last dim %d\"\n k_shape.(q_rank - 1)\n depth;\n let scale_factor =\n match scale with\n | Some s -> s\n | None -> 1.0 /. Stdlib.sqrt (float_of_int depth)\n in\n let transpose_last_two tensor =\n let nd = Array.length (Nx.shape tensor) in\n if nd < 2 then\n invalid_argf_fn \"dot_product_attention\" \"key/value: must have rank >= 2\";\n let axes = Array.init nd Fun.id in\n let tmp = axes.(nd - 1) in\n axes.(nd - 1) <- axes.(nd - 2);\n axes.(nd - 2) <- tmp;\n Nx.transpose tensor ~axes:(Array.to_list axes)\n in\n let k_t = transpose_last_two k in\n let scores = Nx.matmul q k_t in\n let scores =\n if scale_factor = 1.0 then scores\n else Nx.mul scores (Nx.scalar_like scores scale_factor)\n in\n let scores =\n if is_causal then (\n let scores_shape = Nx.shape scores in\n let seq_len_q = scores_shape.(q_rank - 2) in\n let seq_len_k = scores_shape.(q_rank - 1) in\n if seq_len_q <> seq_len_k then\n invalid_argf_fn \"dot_product_attention\"\n \"causal masking requires seq_len_q == seq_len_k\";\n let ones_matrix =\n Nx.full (Nx.dtype scores) [| seq_len_q; seq_len_k |] 1.0\n in\n let causal_mask = Nx.tril ones_matrix in\n let causal_mask = Nx.cast Nx.bool causal_mask in\n let causal_mask = Nx.broadcast_to scores_shape causal_mask in\n let neg_inf = Nx.scalar_like scores (-1e9) in\n Nx.where causal_mask scores neg_inf)\n else scores\n in\n let scores =\n match attention_mask with\n | None -> scores\n | Some mask ->\n let neg_inf = Nx.scalar_like scores (-1e9) in\n Nx.where mask scores neg_inf\n in\n let probs = Nx.softmax ~axes:[ -1 ] scores in\n let probs =\n match dropout_rate with None -> probs | Some rate -> dropout ~rate probs\n in\n Nx.matmul probs v\n\n(* Conv / Pool helpers *)\n\nlet ceildiv a b = (a + b - 1) / b\n\nlet calculate_nn_padding input_spatial ~kernel_size ~stride ~dilation\n ~(padding : [ `Same | `Valid ]) =\n let k = Array.length kernel_size in\n match padding with\n | `Valid -> Array.make k (0, 0)\n | `Same ->\n Array.init k (fun i ->\n let eff_k = (dilation.(i) * (kernel_size.(i) - 1)) + 1 in\n let out = ceildiv input_spatial.(i) stride.(i) in\n let total =\n Stdlib.max 0 (((out - 1) * stride.(i)) + eff_k - input_spatial.(i))\n in\n (total / 2, total - (total / 2)))\n\nlet apply_ceil_mode input_spatial ~kernel_size ~stride ~dilation ~padding\n ~ceil_mode =\n if not ceil_mode then padding\n else\n Array.init (Array.length kernel_size) (fun i ->\n let pb, pa = padding.(i) in\n let padded = input_spatial.(i) + pb + pa in\n let eff_k = (dilation.(i) * (kernel_size.(i) - 1)) + 1 in\n let out_floor = ((padded - eff_k) / stride.(i)) + 1 in\n let out_ceil = ceildiv (padded - eff_k) stride.(i) + 1 in\n if out_ceil > out_floor then\n let extra = ((out_ceil - 1) * stride.(i)) + eff_k - padded in\n (pb, pa + extra)\n else (pb, pa))\n\n(* Convolution *)\n\nlet conv1d ?(groups = 1) ?(stride = 1) ?(dilation = 1) ?(padding = `Valid) ?bias\n x w =\n let x_shape = Nx.shape x in\n let w_shape = Nx.shape w in\n if Array.length x_shape <> 3 then\n invalid_argf_fn \"conv1d\" \"input must be 3D (N, C_in, L)\";\n if Array.length w_shape <> 3 then\n invalid_argf_fn \"conv1d\" \"weight must be 3D (C_out, C_in/groups, K)\";\n let n = x_shape.(0) in\n let cin = x_shape.(1) in\n let cout = w_shape.(0) in\n let cin_per_group = w_shape.(1) in\n if cin <> groups * cin_per_group then\n invalid_argf_fn \"conv1d\" \"C_in=%d does not match groups=%d * C_in/g=%d\" cin\n groups cin_per_group;\n let kernel_size = [| w_shape.(2) |] in\n let stride_arr = [| stride |] in\n let dilation_arr = [| dilation |] in\n let input_spatial = [| x_shape.(2) |] in\n let pad_pairs =\n calculate_nn_padding input_spatial ~kernel_size ~stride:stride_arr\n ~dilation:dilation_arr ~padding\n in\n let kernel_elements = w_shape.(2) in\n (* unfold: (N, C_in, L_in) -> (N, C_in, K, L_out) *)\n let x_unf =\n Nx.extract_patches ~kernel_size ~stride:stride_arr ~dilation:dilation_arr\n ~padding:pad_pairs x\n in\n let x_unf_shape = Nx.shape x_unf in\n let l_out = x_unf_shape.(3) in\n (* Merge channels and kernel: (N, C_in*K, L_out) *)\n let x_col = Nx.reshape [| n; cin * kernel_elements; l_out |] x_unf in\n let result =\n if groups = 1 then\n let w_flat = Nx.reshape [| cout; cin * kernel_elements |] w in\n Nx.matmul w_flat x_col\n else\n let rcout = cout / groups in\n let x_grouped =\n Nx.reshape [| n; groups; cin_per_group * kernel_elements; l_out |] x_col\n in\n let w_grouped =\n Nx.reshape [| groups; rcout; cin_per_group * kernel_elements |] w\n in\n let x_batched =\n Nx.reshape\n [| n * groups; cin_per_group * kernel_elements; l_out |]\n x_grouped\n in\n let w_expanded = Nx.unsqueeze ~axes:[ 0 ] w_grouped in\n let w_expanded =\n Nx.expand\n [| n; groups; rcout; cin_per_group * kernel_elements |]\n w_expanded\n in\n let w_expanded =\n Nx.reshape\n [| n * groups; rcout; cin_per_group * kernel_elements |]\n w_expanded\n in\n let result = Nx.matmul w_expanded x_batched in\n let result = Nx.reshape [| n; groups; rcout; l_out |] result in\n Nx.reshape [| n; cout; l_out |] result\n in\n match bias with\n | None -> result\n | Some b -> Nx.add result (Nx.reshape [| 1; cout; 1 |] b)\n\nlet conv2d ?(groups = 1) ?(stride = (1, 1)) ?(dilation = (1, 1))\n ?(padding = `Valid) ?bias x w =\n let x_shape = Nx.shape x in\n let w_shape = Nx.shape w in\n if Array.length x_shape <> 4 then\n invalid_argf_fn \"conv2d\" \"input must be 4D (N, C_in, H, W)\";\n if Array.length w_shape <> 4 then\n invalid_argf_fn \"conv2d\" \"weight must be 4D (C_out, C_in/groups, kH, kW)\";\n let n = x_shape.(0) in\n let cin = x_shape.(1) in\n let cout = w_shape.(0) in\n let cin_per_group = w_shape.(1) in\n if cin <> groups * cin_per_group then\n invalid_argf_fn \"conv2d\" \"C_in=%d does not match groups=%d * C_in/g=%d\" cin\n groups cin_per_group;\n let sh, sw = stride in\n let dh, dw = dilation in\n let kernel_size = [| w_shape.(2); w_shape.(3) |] in\n let stride_arr = [| sh; sw |] in\n let dilation_arr = [| dh; dw |] in\n let input_spatial = [| x_shape.(2); x_shape.(3) |] in\n let pad_pairs =\n calculate_nn_padding input_spatial ~kernel_size ~stride:stride_arr\n ~dilation:dilation_arr ~padding\n in\n let kernel_elements = w_shape.(2) * w_shape.(3) in\n (* unfold: (N, C_in, H, W) -> (N, C_in, kH*kW, L) *)\n let x_unf =\n Nx.extract_patches ~kernel_size ~stride:stride_arr ~dilation:dilation_arr\n ~padding:pad_pairs x\n in\n let x_unf_shape = Nx.shape x_unf in\n let l_out = x_unf_shape.(3) in\n (* Merge channels and kernel: (N, C_in*kH*kW, L) *)\n let x_col = Nx.reshape [| n; cin * kernel_elements; l_out |] x_unf in\n let result =\n if groups = 1 then\n let w_flat = Nx.reshape [| cout; cin * kernel_elements |] w in\n Nx.matmul w_flat x_col\n else\n let rcout = cout / groups in\n let x_grouped =\n Nx.reshape [| n; groups; cin_per_group * kernel_elements; l_out |] x_col\n in\n let w_grouped =\n Nx.reshape [| groups; rcout; cin_per_group * kernel_elements |] w\n in\n let x_batched =\n Nx.reshape\n [| n * groups; cin_per_group * kernel_elements; l_out |]\n x_grouped\n in\n let w_expanded = Nx.unsqueeze ~axes:[ 0 ] w_grouped in\n let w_expanded =\n Nx.expand\n [| n; groups; rcout; cin_per_group * kernel_elements |]\n w_expanded\n in\n let w_expanded =\n Nx.reshape\n [| n * groups; rcout; cin_per_group * kernel_elements |]\n w_expanded\n in\n let result = Nx.matmul w_expanded x_batched in\n let result = Nx.reshape [| n; groups; rcout; l_out |] result in\n Nx.reshape [| n; cout; l_out |] result\n in\n (* Reshape from (N, C_out, L) to (N, C_out, H_out, W_out) *)\n let padded_h = input_spatial.(0) + fst pad_pairs.(0) + snd pad_pairs.(0) in\n let padded_w = input_spatial.(1) + fst pad_pairs.(1) + snd pad_pairs.(1) in\n let eff_kh = ((kernel_size.(0) - 1) * dh) + 1 in\n let eff_kw = ((kernel_size.(1) - 1) * dw) + 1 in\n let h_out = ((padded_h - eff_kh) / sh) + 1 in\n let w_out = ((padded_w - eff_kw) / sw) + 1 in\n let result = Nx.reshape [| n; cout; h_out; w_out |] result in\n match bias with\n | None -> result\n | Some b -> Nx.add result (Nx.reshape [| 1; cout; 1; 1 |] b)\n\n(* Pooling *)\n\nlet max_pool1d ~kernel_size ?(stride = 1) ?(dilation = 1) ?(padding = `Valid)\n ?(ceil_mode = false) x =\n let x_shape = Nx.shape x in\n if Array.length x_shape <> 3 then\n invalid_argf_fn \"max_pool1d\" \"input must be 3D (N, C, L)\";\n let n = x_shape.(0) in\n let c = x_shape.(1) in\n let kernel_size_arr = [| kernel_size |] in\n let stride_arr = [| stride |] in\n let dilation_arr = [| dilation |] in\n let input_spatial = [| x_shape.(2) |] in\n let pad_pairs =\n calculate_nn_padding input_spatial ~kernel_size:kernel_size_arr\n ~stride:stride_arr ~dilation:dilation_arr ~padding\n in\n let pad_pairs =\n apply_ceil_mode input_spatial ~kernel_size:kernel_size_arr\n ~stride:stride_arr ~dilation:dilation_arr ~padding:pad_pairs ~ceil_mode\n in\n (* unfold: (N, C, L) -> (N, C, K, L_out) *)\n let x_unf =\n Nx.extract_patches ~kernel_size:kernel_size_arr ~stride:stride_arr\n ~dilation:dilation_arr ~padding:pad_pairs x\n in\n let x_unf_ndim = Nx.ndim x_unf in\n let reduced = Nx.max x_unf ~axes:[ x_unf_ndim - 2 ] ~keepdims:false in\n let x_unf_shape = Nx.shape x_unf in\n let l_out = x_unf_shape.(x_unf_ndim - 1) in\n Nx.reshape [| n; c; l_out |] reduced\n\nlet max_pool2d ~kernel_size ?(stride = (1, 1)) ?(dilation = (1, 1))\n ?(padding = `Valid) ?(ceil_mode = false) x =\n let x_shape = Nx.shape x in\n if Array.length x_shape <> 4 then\n invalid_argf_fn \"max_pool2d\" \"input must be 4D (N, C, H, W)\";\n let n = x_shape.(0) in\n let c = x_shape.(1) in\n let kh, kw = kernel_size in\n let sh, sw = stride in\n let dh, dw = dilation in\n let kernel_size_arr = [| kh; kw |] in\n let stride_arr = [| sh; sw |] in\n let dilation_arr = [| dh; dw |] in\n let input_spatial = [| x_shape.(2); x_shape.(3) |] in\n let pad_pairs =\n calculate_nn_padding input_spatial ~kernel_size:kernel_size_arr\n ~stride:stride_arr ~dilation:dilation_arr ~padding\n in\n let pad_pairs =\n apply_ceil_mode input_spatial ~kernel_size:kernel_size_arr\n ~stride:stride_arr ~dilation:dilation_arr ~padding:pad_pairs ~ceil_mode\n in\n let x_unf =\n Nx.extract_patches ~kernel_size:kernel_size_arr ~stride:stride_arr\n ~dilation:dilation_arr ~padding:pad_pairs x\n in\n let x_unf_ndim = Nx.ndim x_unf in\n let reduced = Nx.max x_unf ~axes:[ x_unf_ndim - 2 ] ~keepdims:false in\n let x_unf_shape = Nx.shape x_unf in\n let l_out = x_unf_shape.(x_unf_ndim - 1) in\n let padded_h = input_spatial.(0) + fst pad_pairs.(0) + snd pad_pairs.(0) in\n let padded_w = input_spatial.(1) + fst pad_pairs.(1) + snd pad_pairs.(1) in\n let eff_kh = ((kh - 1) * dh) + 1 in\n let eff_kw = ((kw - 1) * dw) + 1 in\n let h_out = ((padded_h - eff_kh) / sh) + 1 in\n let w_out = ((padded_w - eff_kw) / sw) + 1 in\n let _ = l_out in\n Nx.reshape [| n; c; h_out; w_out |] reduced\n\nlet avg_pool1d ~kernel_size ?(stride = 1) ?(dilation = 1) ?(padding = `Valid)\n ?(ceil_mode = false) ?(count_include_pad = true) x =\n let x_shape = Nx.shape x in\n if Array.length x_shape <> 3 then\n invalid_argf_fn \"avg_pool1d\" \"input must be 3D (N, C, L)\";\n let n = x_shape.(0) in\n let c = x_shape.(1) in\n let kernel_size_arr = [| kernel_size |] in\n let stride_arr = [| stride |] in\n let dilation_arr = [| dilation |] in\n let input_spatial = [| x_shape.(2) |] in\n let pad_pairs =\n calculate_nn_padding input_spatial ~kernel_size:kernel_size_arr\n ~stride:stride_arr ~dilation:dilation_arr ~padding\n in\n let pad_pairs =\n apply_ceil_mode input_spatial ~kernel_size:kernel_size_arr\n ~stride:stride_arr ~dilation:dilation_arr ~padding:pad_pairs ~ceil_mode\n in\n let x_unf =\n Nx.extract_patches ~kernel_size:kernel_size_arr ~stride:stride_arr\n ~dilation:dilation_arr ~padding:pad_pairs x\n in\n let x_unf_ndim = Nx.ndim x_unf in\n let x_unf_shape = Nx.shape x_unf in\n let l_out = x_unf_shape.(x_unf_ndim - 1) in\n let summed = Nx.sum x_unf ~axes:[ x_unf_ndim - 2 ] in\n let result = Nx.reshape [| n; c; l_out |] summed in\n if count_include_pad then Nx.div_s result (float_of_int kernel_size)\n else\n let ones = Nx.ones_like x in\n let ones_unf =\n Nx.extract_patches ~kernel_size:kernel_size_arr ~stride:stride_arr\n ~dilation:dilation_arr ~padding:pad_pairs ones\n in\n let count = Nx.sum ones_unf ~axes:[ Nx.ndim ones_unf - 2 ] in\n let count = Nx.reshape [| n; c; l_out |] count in\n Nx.div result count\n\nlet avg_pool2d ~kernel_size ?(stride = (1, 1)) ?(dilation = (1, 1))\n ?(padding = `Valid) ?(ceil_mode = false) ?(count_include_pad = true) x =\n let x_shape = Nx.shape x in\n if Array.length x_shape <> 4 then\n invalid_argf_fn \"avg_pool2d\" \"input must be 4D (N, C, H, W)\";\n let n = x_shape.(0) in\n let c = x_shape.(1) in\n let kh, kw = kernel_size in\n let sh, sw = stride in\n let dh, dw = dilation in\n let kernel_size_arr = [| kh; kw |] in\n let stride_arr = [| sh; sw |] in\n let dilation_arr = [| dh; dw |] in\n let input_spatial = [| x_shape.(2); x_shape.(3) |] in\n let pad_pairs =\n calculate_nn_padding input_spatial ~kernel_size:kernel_size_arr\n ~stride:stride_arr ~dilation:dilation_arr ~padding\n in\n let pad_pairs =\n apply_ceil_mode input_spatial ~kernel_size:kernel_size_arr\n ~stride:stride_arr ~dilation:dilation_arr ~padding:pad_pairs ~ceil_mode\n in\n let x_unf =\n Nx.extract_patches ~kernel_size:kernel_size_arr ~stride:stride_arr\n ~dilation:dilation_arr ~padding:pad_pairs x\n in\n let x_unf_ndim = Nx.ndim x_unf in\n let x_unf_shape = Nx.shape x_unf in\n let l_out = x_unf_shape.(x_unf_ndim - 1) in\n let summed = Nx.sum x_unf ~axes:[ x_unf_ndim - 2 ] in\n let padded_h = input_spatial.(0) + fst pad_pairs.(0) + snd pad_pairs.(0) in\n let padded_w = input_spatial.(1) + fst pad_pairs.(1) + snd pad_pairs.(1) in\n let eff_kh = ((kh - 1) * dh) + 1 in\n let eff_kw = ((kw - 1) * dw) + 1 in\n let h_out = ((padded_h - eff_kh) / sh) + 1 in\n let w_out = ((padded_w - eff_kw) / sw) + 1 in\n let _ = l_out in\n let result = Nx.reshape [| n; c; h_out; w_out |] summed in\n if count_include_pad then\n let kernel_numel = float_of_int (kh * kw) in\n Nx.div_s result kernel_numel\n else\n let ones = Nx.ones_like x in\n let ones_unf =\n Nx.extract_patches ~kernel_size:kernel_size_arr ~stride:stride_arr\n ~dilation:dilation_arr ~padding:pad_pairs ones\n in\n let count = Nx.sum ones_unf ~axes:[ Nx.ndim ones_unf - 2 ] in\n let count = Nx.reshape [| n; c; h_out; w_out |] count in\n Nx.div result count\n"
},
{
"path": "packages/kaun/lib/fn.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Functional neural network operations.\n\n Stateless building blocks for neural networks: normalization, attention,\n embedding lookup, and regularization. All functions are differentiable\n through Rune's autodiff. *)\n\n(** {1:norm Normalization} *)\n\nval batch_norm :\n ?axes:int list ->\n ?epsilon:float ->\n scale:(float, 'b) Nx.t ->\n bias:(float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t\n(** [batch_norm ?axes ?epsilon ~scale ~bias x] normalizes [x] along [axes], then\n applies learnable [scale] and [bias].\n\n [axes] defaults to [[0]] for 2D and [[0; 2; 3]] for 4D input. [epsilon]\n defaults to [1e-5].\n\n [scale] and [bias] must broadcast across the normalized axes.\n\n Raises [Invalid_argument] if [axes] is empty or out of bounds, or if\n [scale]/[bias] shapes are incompatible. *)\n\nval layer_norm :\n ?axes:int list ->\n ?epsilon:float ->\n ?gamma:(float, 'b) Nx.t ->\n ?beta:(float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t\n(** [layer_norm ?axes ?epsilon ?gamma ?beta x] subtracts the mean and divides by\n the standard deviation along [axes], optionally scaling by [gamma] and\n shifting by [beta].\n\n [axes] defaults to [[-1]]. [epsilon] defaults to [1e-5].\n\n Raises [Invalid_argument] if [axes] is out of bounds, or if [gamma]/[beta]\n shapes are incompatible. *)\n\nval rms_norm :\n ?axes:int list ->\n ?epsilon:float ->\n ?gamma:(float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t\n(** [rms_norm ?axes ?epsilon ?gamma x] normalizes [x] by the root mean square\n along [axes], optionally scaling by [gamma].\n\n [axes] defaults to [[-1]]. [epsilon] defaults to [1e-5].\n\n Raises [Invalid_argument] if [axes] is empty or out of bounds, or if [gamma]\n shape is incompatible. *)\n\n(** {1:embedding Embedding} *)\n\nval embedding :\n ?scale:bool ->\n embedding:(float, 'b) Nx.t ->\n (int32, Nx.int32_elt) Nx.t ->\n (float, 'b) Nx.t\n(** [embedding ?scale ~embedding indices] gathers rows of [embedding] at\n positions given by [indices].\n\n [embedding] must have shape [[vocab_size; embed_dim]]. The result has shape\n [[*indices_shape; embed_dim]].\n\n When [scale] is [true] (the default), the result is multiplied by\n [sqrt(embed_dim)].\n\n Raises [Invalid_argument] if [embedding] is not rank 2 or if [vocab_size] is\n not positive. *)\n\n(** {1:dropout Dropout} *)\n\nval dropout : rate:float -> (float, 'b) Nx.t -> (float, 'b) Nx.t\n(** [dropout ~rate x] randomly zeroes elements of [x] with probability [rate]\n and scales the remaining values by [1 / (1 - rate)].\n\n [rate] must satisfy [0.0 <= rate < 1.0]. When [rate] is [0.0], [x] is\n returned unchanged.\n\n Random keys are drawn from the implicit RNG scope.\n\n Raises [Invalid_argument] if [rate] is out of range or [x] is not floating\n point. *)\n\n(** {1:attention Attention} *)\n\nval dot_product_attention :\n ?attention_mask:(bool, Nx.bool_elt) Nx.t ->\n ?scale:float ->\n ?dropout_rate:float ->\n ?is_causal:bool ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t\n(** [dot_product_attention ?attention_mask ?scale ?dropout_rate ?is_causal q k\n v] is scaled dot-product attention.\n\n [q], [k], [v] must have matching rank (>= 2) and the last dimension of [q]\n and [k] must agree.\n\n [scale] defaults to [1 / sqrt(depth)]. [is_causal] defaults to [false]; when\n [true], a lower-triangular mask is applied (requires\n [seq_len_q = seq_len_k]).\n\n [attention_mask], when provided, broadcasts to the attention score shape:\n [true] keeps scores, [false] sets them to negative infinity.\n\n When [dropout_rate] is set, dropout is applied to attention weights using\n keys from the implicit RNG scope.\n\n Raises [Invalid_argument] if ranks, shapes, or dtypes are incompatible. *)\n\n(** {1:conv Convolution} *)\n\nval conv1d :\n ?groups:int ->\n ?stride:int ->\n ?dilation:int ->\n ?padding:[ `Same | `Valid ] ->\n ?bias:(float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t\n(** [conv1d ?groups ?stride ?dilation ?padding ?bias x w] computes 1D\n convolution.\n\n [x]: [(N, C_in, L)]. [w]: [(C_out, C_in/groups, K)].\n\n [groups] defaults to [1]. [stride] and [dilation] default to [1]. [padding]\n defaults to [`Valid].\n\n Raises [Invalid_argument] if input/weight shapes are incompatible or channel\n counts do not match [groups]. *)\n\nval conv2d :\n ?groups:int ->\n ?stride:int * int ->\n ?dilation:int * int ->\n ?padding:[ `Same | `Valid ] ->\n ?bias:(float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t\n(** [conv2d ?groups ?stride ?dilation ?padding ?bias x w] computes 2D\n convolution.\n\n [x]: [(N, C_in, H, W)]. [w]: [(C_out, C_in/groups, kH, kW)].\n\n [groups] defaults to [1]. [stride] and [dilation] default to [(1, 1)].\n [padding] defaults to [`Valid].\n\n Raises [Invalid_argument] if input/weight shapes are incompatible or channel\n counts do not match [groups]. *)\n\n(** {1:pool Pooling} *)\n\nval max_pool1d :\n kernel_size:int ->\n ?stride:int ->\n ?dilation:int ->\n ?padding:[ `Same | `Valid ] ->\n ?ceil_mode:bool ->\n ('a, 'b) Nx.t ->\n ('a, 'b) Nx.t\n(** [max_pool1d ~kernel_size ?stride ?dilation ?padding ?ceil_mode x] applies 1D\n max pooling.\n\n [x]: [(N, C, L)]. [stride] defaults to [1]. [dilation] defaults to [1].\n [padding] defaults to [`Valid]. [ceil_mode] defaults to [false]. *)\n\nval max_pool2d :\n kernel_size:int * int ->\n ?stride:int * int ->\n ?dilation:int * int ->\n ?padding:[ `Same | `Valid ] ->\n ?ceil_mode:bool ->\n ('a, 'b) Nx.t ->\n ('a, 'b) Nx.t\n(** [max_pool2d ~kernel_size ?stride ?dilation ?padding ?ceil_mode x] applies 2D\n max pooling.\n\n [x]: [(N, C, H, W)]. [stride] defaults to [(1, 1)]. [dilation] defaults to\n [(1, 1)]. [padding] defaults to [`Valid]. [ceil_mode] defaults to [false].\n*)\n\nval avg_pool1d :\n kernel_size:int ->\n ?stride:int ->\n ?dilation:int ->\n ?padding:[ `Same | `Valid ] ->\n ?ceil_mode:bool ->\n ?count_include_pad:bool ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t\n(** [avg_pool1d ~kernel_size ?stride ?dilation ?padding ?ceil_mode\n ?count_include_pad x] applies 1D average pooling.\n\n [x]: [(N, C, L)]. [stride] defaults to [1]. [dilation] defaults to [1].\n [padding] defaults to [`Valid]. [ceil_mode] defaults to [false].\n [count_include_pad] defaults to [true]. *)\n\nval avg_pool2d :\n kernel_size:int * int ->\n ?stride:int * int ->\n ?dilation:int * int ->\n ?padding:[ `Same | `Valid ] ->\n ?ceil_mode:bool ->\n ?count_include_pad:bool ->\n (float, 'b) Nx.t ->\n (float, 'b) Nx.t\n(** [avg_pool2d ~kernel_size ?stride ?dilation ?padding ?ceil_mode\n ?count_include_pad x] applies 2D average pooling.\n\n [x]: [(N, C, H, W)]. [stride] defaults to [(1, 1)]. [dilation] defaults to\n [(1, 1)]. [padding] defaults to [`Valid]. [ceil_mode] defaults to [false].\n [count_include_pad] defaults to [true]. *)\n"
},
{
"path": "packages/kaun/lib/grad.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Dtype = Nx_core.Dtype\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\nlet path_label path = if path = \"\" then \"\" else path\n\nlet err_non_float fn_name path dtype =\n invalid_argf \"%s: %s expected float dtype, got %s\" fn_name (path_label path)\n (Dtype.to_string dtype)\n\nlet err_mismatch fn_name path expected actual =\n invalid_argf \"%s: %s has dtype/layout %s but expected %s\" fn_name\n (path_label path) (Dtype.to_string actual) (Dtype.to_string expected)\n\nlet value_and_grad_aux f params =\n let fn_name = \"Grad.value_and_grad\" in\n let leaves, rebuild = Ptree.flatten params in\n let path_leaves = Ptree.flatten_with_paths params in\n let leaf_count = List.length leaves in\n let path_count = List.length path_leaves in\n if leaf_count <> path_count then\n invalid_argf\n \"%s: internal error: flatten/flatten_with_paths length mismatch (%d vs \\\n %d)\"\n fn_name leaf_count path_count;\n if leaf_count = 0 then\n let value, aux = f params in\n (value, rebuild [], aux)\n else\n let leaves_array = Array.of_list leaves in\n let paths_array = Array.of_list (List.map fst path_leaves) in\n let first_leaf = leaves_array.(0) in\n let first_path = paths_array.(0) in\n Ptree.with_tensor first_leaf\n {\n run =\n (fun (type a layout) (first_tensor : (a, layout) Nx.t) ->\n let first_dtype : (a, layout) Dtype.t = Nx.dtype first_tensor in\n if not (Dtype.is_float first_dtype) then\n err_non_float fn_name first_path first_dtype;\n let typed_inputs =\n List.mapi\n (fun index leaf ->\n let path = paths_array.(index) in\n Ptree.with_tensor leaf\n {\n run =\n (fun (type a2 layout2) (tensor : (a2, layout2) Nx.t) ->\n let dtype = Nx.dtype tensor in\n if not (Dtype.is_float dtype) then\n err_non_float fn_name path dtype;\n match Dtype.equal_witness first_dtype dtype with\n | Some Type.Equal -> (tensor : (a, layout) Nx.t)\n | None -> err_mismatch fn_name path first_dtype dtype);\n })\n leaves\n in\n let aux = ref None in\n let objective typed_params =\n let packed =\n List.map (fun tensor -> Ptree.P tensor) typed_params\n in\n let value, aux_value = f (rebuild packed) in\n if Option.is_none !aux then aux := Some aux_value;\n value\n in\n let value, grads = Rune.value_and_grads objective typed_inputs in\n let aux_value =\n match !aux with\n | Some value -> value\n | None ->\n invalid_argf\n \"%s: internal error: objective did not produce auxiliary \\\n output\"\n fn_name\n in\n let grad_leaves = List.map (fun grad -> Ptree.P grad) grads in\n (value, rebuild grad_leaves, aux_value));\n }\n\nlet value_and_grad f params =\n let value, grads, () = value_and_grad_aux (fun tree -> (f tree, ())) params in\n (value, grads)\n\nlet value_and_grad_mixed f params =\n let fn_name = \"Grad.value_and_grad_mixed\" in\n let leaves, rebuild = Ptree.flatten params in\n if List.length leaves = 0 then (f params, rebuild [])\n else\n let path_leaves = Ptree.flatten_with_paths params in\n let leaf_count = List.length leaves in\n let path_count = List.length path_leaves in\n if leaf_count <> path_count then\n invalid_argf\n \"%s: internal error: flatten/flatten_with_paths length mismatch (%d vs \\\n %d)\"\n fn_name leaf_count path_count;\n let leaves_array = Array.of_list leaves in\n let paths_array = Array.of_list (List.map fst path_leaves) in\n let grads_array = Array.make leaf_count None in\n let groups = Hashtbl.create 8 in\n Array.iteri\n (fun index (Ptree.P tensor) ->\n let dtype = Nx.dtype tensor in\n if not (Dtype.is_float dtype) then\n err_non_float fn_name paths_array.(index) dtype;\n let group_key = Dtype.to_string dtype in\n match Hashtbl.find_opt groups group_key with\n | None -> Hashtbl.add groups group_key (Ptree.P tensor, [ index ])\n | Some (repr, indices) ->\n Hashtbl.replace groups group_key (repr, index :: indices))\n leaves_array;\n let grouped_indices =\n Hashtbl.fold\n (fun _ (repr, indices) acc -> (repr, List.rev indices) :: acc)\n groups []\n in\n let value = ref None in\n List.iter\n (fun (repr, indices) ->\n Ptree.with_tensor repr\n {\n run =\n (fun (type a layout) (repr_tensor : (a, layout) Nx.t) ->\n let repr_dtype : (a, layout) Dtype.t = Nx.dtype repr_tensor in\n let typed_inputs =\n List.map\n (fun index ->\n Ptree.with_tensor leaves_array.(index)\n {\n run =\n (fun (type a2 layout2)\n (tensor : (a2, layout2) Nx.t)\n ->\n let dtype = Nx.dtype tensor in\n match Dtype.equal_witness repr_dtype dtype with\n | Some Type.Equal -> (tensor : (a, layout) Nx.t)\n | None ->\n err_mismatch fn_name paths_array.(index)\n repr_dtype dtype);\n })\n indices\n in\n let objective group_params =\n let packed = Array.copy leaves_array in\n List.iter2\n (fun index tensor -> packed.(index) <- Ptree.P tensor)\n indices group_params;\n f (rebuild (Array.to_list packed))\n in\n let current_value, current_grads =\n Rune.value_and_grads objective typed_inputs\n in\n if Option.is_none !value then value := Some current_value;\n List.iter2\n (fun index grad -> grads_array.(index) <- Some (Ptree.P grad))\n indices current_grads);\n })\n grouped_indices;\n let value =\n match !value with\n | Some v -> v\n | None ->\n invalid_argf \"%s: internal error: no autodiff group produced a value\"\n fn_name\n in\n let grad_leaves =\n Array.to_list\n (Array.mapi\n (fun index grad ->\n match grad with\n | Some g -> g\n | None ->\n invalid_argf \"%s: internal error: missing gradient for leaf %s\"\n fn_name\n (path_label paths_array.(index)))\n grads_array)\n in\n (value, rebuild grad_leaves)\n\nlet grad f params = snd (value_and_grad f params)\n"
},
{
"path": "packages/kaun/lib/grad.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Automatic differentiation over parameter trees.\n\n {!value_and_grad} differentiates scalar losses with respect to {!Ptree.t}\n leaves. By default, all leaves must share one floating dtype; this enables a\n single forward/backward pass.\n\n Use {!value_and_grad_aux} to return auxiliary data (for example updated\n layer state) alongside the loss. Use {!value_and_grad_mixed} when mixed\n dtypes are required. *)\n\n(** {1:core Core} *)\n\nval value_and_grad :\n (Ptree.t -> (float, 'l) Nx.t) -> Ptree.t -> (float, 'l) Nx.t * Ptree.t\n(** [value_and_grad f params] is [(f params, grads)].\n\n [params] must contain only floating-point leaves and all leaves must have\n the same dtype/layout witness.\n\n Raises [Invalid_argument] if a leaf is non-float, or if dtypes/layout differ\n across leaves. Error messages include leaf paths. *)\n\nval value_and_grad_aux :\n (Ptree.t -> (float, 'l) Nx.t * 'aux) ->\n Ptree.t ->\n (float, 'l) Nx.t * Ptree.t * 'aux\n(** [value_and_grad_aux f params] differentiates [fst (f params)] and returns\n [(loss, grads, aux)].\n\n The same dtype constraints and errors as {!value_and_grad} apply. *)\n\nval value_and_grad_mixed :\n (Ptree.t -> (float, 'l) Nx.t) -> Ptree.t -> (float, 'l) Nx.t * Ptree.t\n(** [value_and_grad_mixed f params] supports mixed floating dtypes/layouts by\n grouping leaves and running multiple autodiff passes.\n\n {b Warning.} [f] may be evaluated multiple times (once per dtype/layout\n group).\n\n Raises [Invalid_argument] if any leaf is non-float. Error messages include\n leaf paths. *)\n\nval grad : (Ptree.t -> (float, 'l) Nx.t) -> Ptree.t -> Ptree.t\n(** [grad f params] is [snd (value_and_grad f params)]. *)\n"
},
{
"path": "packages/kaun/lib/hf/dune",
"content": "(library\n (name kaun_hf)\n (public_name kaun.hf)\n (libraries unix rune nx nx.io kaun jsont jsont.bytesrw))\n"
},
{
"path": "packages/kaun/lib/hf/kaun_hf.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Types *)\n\ntype revision = Main | Rev of string\n\n(* Error messages *)\n\nlet err_no_curl = \"curl not found on PATH\"\nlet err_download url = Printf.sprintf \"Failed to download %s\" url\n\nlet err_offline model_id filename =\n Printf.sprintf \"Not cached (offline): %s/%s\" model_id filename\n\nlet err_no_safetensors model_id =\n Printf.sprintf \"No safetensors found for %s\" model_id\n\nlet err_missing_tensor model_id name path =\n Printf.sprintf \"%s: tensor %S missing in shard %s\" model_id name path\n\nlet err_empty_weight_map = \"Empty weight_map in index file\"\nlet err_missing_weight_map = \"Missing weight_map in index file\"\n\nlet err_incomplete_shards =\n \"Incomplete shard loading: not all weight_map tensors were found\"\n\n(* Cache directory *)\n\nlet default_cache_dir () =\n match Sys.getenv_opt \"RAVEN_CACHE_ROOT\" with\n | Some d when d <> \"\" -> Filename.concat d \"huggingface\"\n | _ ->\n let xdg =\n match Sys.getenv_opt \"XDG_CACHE_HOME\" with\n | Some d when d <> \"\" -> d\n | _ -> Filename.concat (Sys.getenv \"HOME\") \".cache\"\n in\n Filename.concat (Filename.concat xdg \"raven\") \"huggingface\"\n\n(* Filesystem *)\n\nlet rec mkdir_p path =\n if path = \"\" || path = \".\" || path = Filename.dir_sep then ()\n else if not (Sys.file_exists path) then begin\n mkdir_p (Filename.dirname path);\n try Unix.mkdir path 0o755 with Unix.Unix_error (Unix.EEXIST, _, _) -> ()\n end\n\nlet rec rm_rf path =\n if Sys.is_directory path then begin\n let entries = Sys.readdir path in\n Array.iter (fun e -> rm_rf (Filename.concat path e)) entries;\n Unix.rmdir path\n end\n else Sys.remove path\n\n(* HTTP via curl *)\n\nlet curl_available =\n lazy (Unix.system \"command -v curl >/dev/null 2>&1\" = Unix.WEXITED 0)\n\nlet check_curl () = if not (Lazy.force curl_available) then failwith err_no_curl\n\nlet header_flags headers =\n List.map\n (fun (k, v) -> Printf.sprintf \"-H %s\" (Filename.quote (k ^ \": \" ^ v)))\n headers\n |> String.concat \" \"\n\nlet curl_download ~headers ~url ~dest () =\n check_curl ();\n mkdir_p (Filename.dirname dest);\n let hdr = header_flags headers in\n let cmd =\n Printf.sprintf \"curl -L --fail -s %s -o %s %s\" hdr (Filename.quote dest)\n (Filename.quote url)\n in\n match Unix.system cmd with\n | Unix.WEXITED 0 -> ()\n | _ ->\n (try Sys.remove dest with Sys_error _ -> ());\n failwith (err_download url)\n\n(* Hub URL and cache paths *)\n\nlet revision_string = function Main -> \"main\" | Rev r -> r\n\nlet hub_url ~model_id ~revision ~filename =\n Printf.sprintf \"https://huggingface.co/%s/resolve/%s/%s\" model_id\n (revision_string revision) filename\n\nlet sanitize_model_id model_id =\n String.map (fun c -> if c = '/' then '-' else c) model_id\n\nlet cache_path ~cache_dir ~model_id ~revision ~filename =\n let rev = revision_string revision in\n let model_dir = sanitize_model_id model_id in\n Filename.concat cache_dir\n (Filename.concat model_dir (Filename.concat rev filename))\n\nlet auth_headers = function\n | Some t -> [ (\"Authorization\", \"Bearer \" ^ t) ]\n | None -> []\n\n(* Downloading *)\n\nlet download_file ?token ?cache_dir ?(offline = false) ?(revision = Main)\n ~model_id ~filename () =\n let token =\n match token with Some _ as t -> t | None -> Sys.getenv_opt \"HF_TOKEN\"\n in\n let cache_dir = Option.value cache_dir ~default:(default_cache_dir ()) in\n let local = cache_path ~cache_dir ~model_id ~revision ~filename in\n if Sys.file_exists local then local\n else if offline then failwith (err_offline model_id filename)\n else begin\n let url = hub_url ~model_id ~revision ~filename in\n curl_download ~headers:(auth_headers token) ~url ~dest:local ();\n local\n end\n\n(* JSON helpers *)\n\nlet read_json_file path =\n let ic = open_in path in\n let s =\n Fun.protect\n ~finally:(fun () -> close_in ic)\n (fun () -> really_input_string ic (in_channel_length ic))\n in\n match Jsont_bytesrw.decode_string Jsont.json s with\n | Ok v -> v\n | Error e -> failwith e\n\nlet json_mem name = function\n | Jsont.Object (mems, _) -> (\n match Jsont.Json.find_mem name mems with\n | Some (_, v) -> v\n | None -> Jsont.Null ((), Jsont.Meta.none))\n | _ -> Jsont.Null ((), Jsont.Meta.none)\n\n(* Tensor conversion *)\n\nlet to_ptree_tensor (Nx_io.P nx) = Kaun.Ptree.P nx\n\n(* Loading *)\n\nlet load_entries ?allowed_names path =\n let archive = Nx_io.load_safetensors path in\n match allowed_names with\n | None ->\n Hashtbl.fold\n (fun name packed acc -> (name, to_ptree_tensor packed) :: acc)\n archive []\n | Some names ->\n List.map\n (fun name ->\n match Hashtbl.find_opt archive name with\n | Some packed -> (name, to_ptree_tensor packed)\n | None -> failwith (err_missing_tensor \"\" name path))\n names\n\nlet try_download f =\n try Some (f ()) with Failure _ -> None | Sys_error _ -> None\n\nlet load_sharded ~download index_filename =\n match try_download (fun () -> download index_filename) with\n | None -> None\n | Some index_path ->\n let json = read_json_file index_path in\n let weight_map =\n match json_mem \"weight_map\" json with\n | Jsont.Object (entries, _) ->\n List.map\n (fun ((tensor_name, _), shard_json) ->\n match shard_json with\n | Jsont.String (shard, _) -> (tensor_name, shard)\n | _ -> failwith err_missing_weight_map)\n entries\n | _ -> failwith err_missing_weight_map\n in\n if weight_map = [] then failwith err_empty_weight_map;\n (* Group tensors by shard filename, preserving file order *)\n let shards_by_file = Hashtbl.create 8 in\n let file_order = ref [] in\n List.iter\n (fun (tensor_name, shard_filename) ->\n match Hashtbl.find_opt shards_by_file shard_filename with\n | Some tensors ->\n Hashtbl.replace shards_by_file shard_filename\n (tensor_name :: tensors)\n | None ->\n Hashtbl.add shards_by_file shard_filename [ tensor_name ];\n file_order := shard_filename :: !file_order)\n weight_map;\n let file_order = List.rev !file_order in\n let seen = Hashtbl.create (List.length weight_map) in\n let entries =\n List.fold_left\n (fun acc shard_filename ->\n let shard_path = download shard_filename in\n let tensors =\n match Hashtbl.find_opt shards_by_file shard_filename with\n | Some names -> List.rev names\n | None -> []\n in\n let new_entries = load_entries ~allowed_names:tensors shard_path in\n List.iter\n (fun (name, _) -> Hashtbl.replace seen name ())\n new_entries;\n List.rev_append new_entries acc)\n [] file_order\n in\n if Hashtbl.length seen <> List.length weight_map then\n failwith err_incomplete_shards;\n Some (List.rev entries)\n\nlet load_single ~download filename =\n match try_download (fun () -> download filename) with\n | None -> None\n | Some path -> Some (load_entries path)\n\nlet load_config ?token ?cache_dir ?offline ?revision ~model_id () =\n let path =\n download_file ?token ?cache_dir ?offline ?revision ~model_id\n ~filename:\"config.json\" ()\n in\n read_json_file path\n\nlet load_weights ?token ?cache_dir ?offline ?revision ~model_id () =\n let download filename =\n download_file ?token ?cache_dir ?offline ?revision ~model_id ~filename ()\n in\n match load_sharded ~download \"model.safetensors.index.json\" with\n | Some entries -> entries\n | None -> (\n match load_single ~download \"model.safetensors\" with\n | Some entries -> entries\n | None -> failwith (err_no_safetensors model_id))\n\n(* Cache management *)\n\nlet clear_cache ?cache_dir ?model_id () =\n let cache_dir = Option.value cache_dir ~default:(default_cache_dir ()) in\n match model_id with\n | Some id ->\n let path = Filename.concat cache_dir (sanitize_model_id id) in\n if Sys.file_exists path then rm_rf path\n | None -> if Sys.file_exists cache_dir then rm_rf cache_dir\n"
},
{
"path": "packages/kaun/lib/hf/kaun_hf.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** HuggingFace Hub integration.\n\n Download pretrained model weights and configuration files from the\n {{:https://huggingface.co}HuggingFace Hub}. Supports single-file and sharded\n SafeTensors checkpoints, caching, authentication, and offline mode.\n\n {[\n let config =\n Kaun_hf.load_config ~model_id:\"bert-base-uncased\" ()\n in\n let weights =\n Kaun_hf.load_weights ~model_id:\"bert-base-uncased\" ()\n in\n (* weights : (string * Kaun.Ptree.tensor) list *)\n ]} *)\n\n(** {1:types Types} *)\n\n(** The type for repository revisions. *)\ntype revision =\n | Main (** The default branch. *)\n | Rev of string (** A tag, branch name, or commit hash. *)\n\n(** {1:downloading Downloading} *)\n\nval download_file :\n ?token:string ->\n ?cache_dir:string ->\n ?offline:bool ->\n ?revision:revision ->\n model_id:string ->\n filename:string ->\n unit ->\n string\n(** [download_file ~model_id ~filename ()] is the local path to [filename] from\n the repository [model_id].\n\n The file is downloaded to the cache on first access and served from there on\n subsequent calls.\n\n [token] is a HuggingFace API token for private repositories. Defaults to the\n value of [HF_TOKEN].\n\n [cache_dir] defaults to [{RAVEN_CACHE_ROOT}/huggingface], or\n [{XDG_CACHE_HOME}/raven/huggingface] when unset.\n\n [offline] defaults to [false]. When [true], only cached files are returned.\n\n [revision] defaults to {!Main}.\n\n Raises [Failure] if the download fails or the file is not cached in offline\n mode. *)\n\n(** {1:loading Loading} *)\n\nval load_config :\n ?token:string ->\n ?cache_dir:string ->\n ?offline:bool ->\n ?revision:revision ->\n model_id:string ->\n unit ->\n Jsont.json\n(** [load_config ~model_id ()] is the parsed [config.json] from [model_id].\n\n Parameters are the same as {!download_file}.\n\n Raises [Failure] on download or JSON parse errors. *)\n\nval load_weights :\n ?token:string ->\n ?cache_dir:string ->\n ?offline:bool ->\n ?revision:revision ->\n model_id:string ->\n unit ->\n (string * Kaun.Ptree.tensor) list\n(** [load_weights ~model_id ()] is the list of [(name, tensor)] pairs from\n [model_id]'s SafeTensors checkpoint.\n\n Handles sharded checkpoints transparently: when\n [model.safetensors.index.json] is present, all referenced shards are\n downloaded and merged. Falls back to [model.safetensors] when no index\n exists.\n\n Tensor names are the raw keys from the SafeTensors file (e.g.\n [\"bert.encoder.layer.0.attention.self.query.weight\"]). Model code is\n responsible for mapping these to its own parameter structure.\n\n Parameters are the same as {!download_file}.\n\n Raises [Failure] if no SafeTensors files are found, or on download/parse\n errors. *)\n\n(** {1:cache Cache management} *)\n\nval clear_cache : ?cache_dir:string -> ?model_id:string -> unit -> unit\n(** [clear_cache ()] removes all cached files.\n\n When [model_id] is given, only that model's cache is removed. *)\n"
},
{
"path": "packages/kaun/lib/init.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype t = {\n f : 'layout. int array -> (float, 'layout) Nx.dtype -> (float, 'layout) Nx.t;\n}\n\ntype mode = [ `Fan_in | `Fan_out | `Fan_avg ]\ntype distribution = [ `Normal | `Truncated_normal | `Uniform ]\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\n\nlet check_non_negative what value =\n if value < 0.0 then invalid_argf \"%s must be >= 0, got %g\" what value\n\nlet normalize_axis ~rank ~name axis =\n let axis = if axis < 0 then rank + axis else axis in\n if axis < 0 || axis >= rank then\n invalid_argf \"invalid %s axis: %d for rank-%d shape\" name axis rank;\n axis\n\n(* Fan computation for variance scaling. *)\n\nlet compute_fans shape ~in_axis ~out_axis =\n let rank = Array.length shape in\n if rank = 0 then (1, 1)\n else if rank = 1 then\n let total = shape.(0) in\n (total, total)\n else\n let in_axis = normalize_axis ~rank ~name:\"in\" in_axis in\n let out_axis = normalize_axis ~rank ~name:\"out\" out_axis in\n let fan_in = shape.(in_axis) in\n let fan_out = shape.(out_axis) in\n let receptive = ref 1 in\n for i = 0 to rank - 1 do\n if i <> in_axis && i <> out_axis then receptive := !receptive * shape.(i)\n done;\n (fan_in * !receptive, fan_out * !receptive)\n\n(* Truncated normal with bounds at +/-2 standard deviations. *)\n\nlet truncated_normal ~stddev shape dtype =\n let z = Nx.truncated_normal dtype ~lower:(-2.0) ~upper:2.0 shape in\n Nx.mul z (Nx.scalar dtype stddev)\n\n(* Variance scaling — the general framework behind glorot/he/lecun. *)\n\nlet variance_scaling ~scale ~mode ~distribution ?(in_axis = -2) ?(out_axis = -1)\n () =\n check_non_negative \"scale\" scale;\n {\n f =\n (fun shape dtype ->\n let fan_in, fan_out = compute_fans shape ~in_axis ~out_axis in\n let n =\n match mode with\n | `Fan_in -> float_of_int fan_in\n | `Fan_out -> float_of_int fan_out\n | `Fan_avg -> float_of_int (fan_in + fan_out) /. 2.0\n in\n if n <= 0.0 then\n invalid_argf \"non-positive fan: fan_in=%d fan_out=%d\" fan_in fan_out;\n let variance = scale /. n in\n match distribution with\n | `Normal ->\n let z = Nx.randn dtype shape in\n Nx.mul z (Nx.scalar dtype (sqrt variance))\n | `Truncated_normal ->\n (* Correct for stddev loss from truncation to [-2, 2]. *)\n truncated_normal\n ~stddev:(sqrt variance /. 0.87962566103423978)\n shape dtype\n | `Uniform ->\n let limit = sqrt (3.0 *. variance) in\n let u = Nx.rand dtype shape in\n Nx.sub\n (Nx.mul u (Nx.scalar dtype (2.0 *. limit)))\n (Nx.scalar dtype limit));\n }\n\n(* Constant *)\n\nlet constant value = { f = (fun shape dtype -> Nx.full dtype shape value) }\nlet zeros = constant 0.0\nlet ones = constant 1.0\n\n(* Random *)\n\nlet uniform ?(scale = 0.01) () =\n check_non_negative \"scale\" scale;\n {\n f =\n (fun shape dtype -> Nx.mul (Nx.rand dtype shape) (Nx.scalar dtype scale));\n }\n\nlet normal ?(stddev = 0.01) () =\n check_non_negative \"stddev\" stddev;\n {\n f =\n (fun shape dtype ->\n Nx.mul (Nx.randn dtype shape) (Nx.scalar dtype stddev));\n }\n\n(* Glorot / Xavier *)\n\nlet glorot_uniform ?(in_axis = -2) ?(out_axis = -1) () =\n variance_scaling ~scale:1.0 ~mode:`Fan_avg ~distribution:`Uniform ~in_axis\n ~out_axis ()\n\nlet glorot_normal ?(in_axis = -2) ?(out_axis = -1) () =\n variance_scaling ~scale:1.0 ~mode:`Fan_avg ~distribution:`Truncated_normal\n ~in_axis ~out_axis ()\n\n(* He / Kaiming *)\n\nlet he_uniform ?(in_axis = -2) ?(out_axis = -1) () =\n variance_scaling ~scale:2.0 ~mode:`Fan_in ~distribution:`Uniform ~in_axis\n ~out_axis ()\n\nlet he_normal ?(in_axis = -2) ?(out_axis = -1) () =\n variance_scaling ~scale:2.0 ~mode:`Fan_in ~distribution:`Truncated_normal\n ~in_axis ~out_axis ()\n\n(* LeCun *)\n\nlet lecun_uniform ?(in_axis = -2) ?(out_axis = -1) () =\n variance_scaling ~scale:1.0 ~mode:`Fan_in ~distribution:`Uniform ~in_axis\n ~out_axis ()\n\nlet lecun_normal ?(in_axis = -2) ?(out_axis = -1) () =\n variance_scaling ~scale:1.0 ~mode:`Fan_in ~distribution:`Truncated_normal\n ~in_axis ~out_axis ()\n"
},
{
"path": "packages/kaun/lib/init.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Weight initialization strategies.\n\n Initializers map a shape and float dtype to tensors. Random keys are\n obtained implicitly via {!Nx.Rng.next_key}. Named families (Glorot, He,\n LeCun) are defined in terms of {!variance_scaling}. *)\n\n(** {1:types Types} *)\n\ntype t = {\n f : 'layout. int array -> (float, 'layout) Nx.dtype -> (float, 'layout) Nx.t;\n}\n(** [t] is the type for initializers.\n\n [i.f shape dtype] is an initialized tensor for [shape] and [dtype]. Random\n keys are drawn from the implicit RNG scope. *)\n\n(** {1:constant Constant} *)\n\nval zeros : t\n(** [zeros] is the initializer that fills with [0.0]. *)\n\nval ones : t\n(** [ones] is the initializer that fills with [1.0]. *)\n\nval constant : float -> t\n(** [constant v] is the initializer that fills with [v]. *)\n\n(** {1:random Random} *)\n\nval uniform : ?scale:float -> unit -> t\n(** [uniform ?scale ()] is the initializer that samples from [U(0, scale)].\n\n [scale] defaults to [0.01].\n\n Raises [Invalid_argument] if [scale] is negative. *)\n\nval normal : ?stddev:float -> unit -> t\n(** [normal ?stddev ()] is the initializer that samples from [N(0, stddev)].\n\n [stddev] defaults to [0.01].\n\n Raises [Invalid_argument] if [stddev] is negative. *)\n\n(** {1:variance Variance Scaling} *)\n\ntype mode = [ `Fan_in | `Fan_out | `Fan_avg ]\n(** The type for variance-scaling fan modes. *)\n\ntype distribution = [ `Normal | `Truncated_normal | `Uniform ]\n(** The type for variance-scaling distribution families. *)\n\nval variance_scaling :\n scale:float ->\n mode:mode ->\n distribution:distribution ->\n ?in_axis:int ->\n ?out_axis:int ->\n unit ->\n t\n(** [variance_scaling ~scale ~mode ~distribution ?in_axis ?out_axis ()] is the\n variance-scaling initializer.\n\n [in_axis] defaults to [-2] and [out_axis] defaults to [-1]. Negative axes\n are interpreted from the end.\n\n The target variance is [scale / n], with:\n - [n = fan_in] for [`Fan_in].\n - [n = fan_out] for [`Fan_out].\n - [n = (fan_in + fan_out) / 2] for [`Fan_avg].\n\n Distributions are:\n - [`Normal]: [N(0, scale / n)].\n - [`Uniform]: [U(-limit, limit)] with [limit = sqrt (3 * scale / n)].\n - [`Truncated_normal]: normal samples truncated to \\[[-2];[2]\\] and rescaled\n to match [scale / n].\n\n Raises [Invalid_argument] if:\n - [scale] is negative.\n - [in_axis] or [out_axis] is out of bounds for rank > 1.\n - the computed fan is non-positive. *)\n\n(** {1:glorot Glorot/Xavier} *)\n\nval glorot_uniform : ?in_axis:int -> ?out_axis:int -> unit -> t\n(** [glorot_uniform ?in_axis ?out_axis ()] is Glorot/Xavier uniform\n initialization.\n\n It samples from [U(-limit, limit)] with\n [limit = sqrt (6 / (fan_in + fan_out))].\n\n This is the Xavier/Glorot scheme of Glorot and Bengio (2010). It is\n implemented via {!variance_scaling} with fan-average mode.\n\n Raises [Invalid_argument] under the same conditions as {!variance_scaling}.\n*)\n\nval glorot_normal : ?in_axis:int -> ?out_axis:int -> unit -> t\n(** [glorot_normal ?in_axis ?out_axis ()] is Glorot/Xavier normal\n initialization.\n\n It uses truncated normal sampling with fan-average target variance\n [2 / (fan_in + fan_out)].\n\n This is the Xavier/Glorot family of Glorot and Bengio (2010). It is\n implemented via {!variance_scaling}.\n\n Raises [Invalid_argument] under the same conditions as {!variance_scaling}.\n*)\n\n(** {1:he He/Kaiming} *)\n\nval he_uniform : ?in_axis:int -> ?out_axis:int -> unit -> t\n(** [he_uniform ?in_axis ?out_axis ()] is He/Kaiming uniform initialization.\n\n It samples from [U(-limit, limit)] with [limit = sqrt (6 / fan_in)].\n\n This is the Kaiming/He scheme of He et al. (2015), commonly used for\n ReLU-like activations. It is implemented via {!variance_scaling} in fan-in\n mode.\n\n Raises [Invalid_argument] under the same conditions as {!variance_scaling}.\n*)\n\nval he_normal : ?in_axis:int -> ?out_axis:int -> unit -> t\n(** [he_normal ?in_axis ?out_axis ()] is He/Kaiming normal initialization.\n\n It uses truncated normal sampling with fan-in target variance [2 / fan_in].\n\n This is the Kaiming/He family of He et al. (2015). It is implemented via\n {!variance_scaling}.\n\n Raises [Invalid_argument] under the same conditions as {!variance_scaling}.\n*)\n\n(** {1:lecun LeCun} *)\n\nval lecun_uniform : ?in_axis:int -> ?out_axis:int -> unit -> t\n(** [lecun_uniform ?in_axis ?out_axis ()] is LeCun uniform initialization.\n\n It samples from [U(-limit, limit)] with [limit = sqrt (3 / fan_in)].\n\n This is the LeCun fan-in family (Efficient BackProp, LeCun et al., 1998). It\n is implemented via {!variance_scaling}.\n\n Raises [Invalid_argument] under the same conditions as {!variance_scaling}.\n*)\n\nval lecun_normal : ?in_axis:int -> ?out_axis:int -> unit -> t\n(** [lecun_normal ?in_axis ?out_axis ()] is LeCun normal initialization.\n\n It uses truncated normal sampling with fan-in target variance [1 / fan_in].\n\n This is the LeCun fan-in family (Efficient BackProp, LeCun et al., 1998). It\n is implemented via {!variance_scaling}.\n\n Raises [Invalid_argument] under the same conditions as {!variance_scaling}.\n*)\n"
},
{
"path": "packages/kaun/lib/layer.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype 'layout vars = {\n params : Ptree.t;\n state : Ptree.t;\n dtype : (float, 'layout) Nx.dtype;\n}\n\ntype ('input, 'output) t = {\n init : 'layout. dtype:(float, 'layout) Nx.dtype -> 'layout vars;\n apply :\n 'layout 'in_elt.\n params:Ptree.t ->\n state:Ptree.t ->\n dtype:(float, 'layout) Nx.dtype ->\n training:bool ->\n ?ctx:Context.t ->\n ('input, 'in_elt) Nx.t ->\n ('output, 'layout) Nx.t * Ptree.t;\n}\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\nlet params v = v.params\nlet state v = v.state\nlet dtype v = v.dtype\nlet with_params v params = { v with params }\nlet with_state v state = { v with state }\nlet make_vars ~params ~state ~dtype = { params; state; dtype }\n\nmodule Dtype = Nx_core.Dtype\n\nlet require_same_float_dtype (type p in_elt) ~ctx\n (expected : (float, p) Nx.dtype) (x : (float, in_elt) Nx.t) :\n (float, p) Nx.t =\n match Dtype.equal_witness expected (Nx.dtype x) with\n | Some Type.Equal -> (x : (float, p) Nx.t)\n | None ->\n invalid_argf \"%s: input dtype %s does not match model dtype %s\" ctx\n (Dtype.to_string (Nx.dtype x))\n (Dtype.to_string expected)\n\nlet require_int32_indices (type in_elt) ~ctx (x : (int32, in_elt) Nx.t) :\n (int32, Bigarray.int32_elt) Nx.t =\n match Dtype.equal_witness Nx.int32 (Nx.dtype x) with\n | Some Type.Equal -> (x : (int32, Bigarray.int32_elt) Nx.t)\n | None ->\n invalid_argf \"%s: expected int32 indices, got %s\" ctx\n (Dtype.to_string (Nx.dtype x))\n\nlet init t ~dtype = t.init ~dtype\n\nlet apply (type a b layout in_elt) (t : (a, b) t) (vars : layout vars) ~training\n ?ctx (x : (a, in_elt) Nx.t) =\n let y, state =\n t.apply ~params:vars.params ~state:vars.state ~dtype:vars.dtype ~training\n ?ctx x\n in\n (y, { vars with state })\n\nlet compose left right =\n {\n init =\n (fun ~dtype ->\n let k1 = Nx.Rng.next_key () in\n let k2 = Nx.Rng.next_key () in\n let left_vars = Nx.Rng.with_key k1 (fun () -> left.init ~dtype) in\n let right_vars = Nx.Rng.with_key k2 (fun () -> right.init ~dtype) in\n {\n params =\n Ptree.dict\n [ (\"left\", left_vars.params); (\"right\", right_vars.params) ];\n state =\n Ptree.dict\n [ (\"left\", left_vars.state); (\"right\", right_vars.state) ];\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n let param_fields =\n Ptree.Dict.fields_exn ~ctx:\"Layer.compose.params\" params\n in\n let state_fields =\n Ptree.Dict.fields_exn ~ctx:\"Layer.compose.state\" state\n in\n let left_params =\n Ptree.Dict.find_exn ~ctx:\"Layer.compose.params\" \"left\" param_fields\n in\n let right_params =\n Ptree.Dict.find_exn ~ctx:\"Layer.compose.params\" \"right\" param_fields\n in\n let left_state =\n Ptree.Dict.find_exn ~ctx:\"Layer.compose.state\" \"left\" state_fields\n in\n let right_state =\n Ptree.Dict.find_exn ~ctx:\"Layer.compose.state\" \"right\" state_fields\n in\n let y, left_state' =\n left.apply ~params:left_params ~state:left_state ~dtype ~training ?ctx\n x\n in\n let z, right_state' =\n right.apply ~params:right_params ~state:right_state ~dtype ~training\n ?ctx y\n in\n (z, Ptree.dict [ (\"left\", left_state'); (\"right\", right_state') ]));\n }\n\n(* Dense *)\n\nlet linear ~in_features ~out_features ?weight_init ?bias_init () =\n let weight_init =\n match weight_init with\n | Some init_value -> init_value\n | None -> Init.glorot_uniform ()\n in\n let bias_init =\n match bias_init with Some init_value -> init_value | None -> Init.zeros\n in\n {\n init =\n (fun ~dtype ->\n let weight = weight_init.f [| in_features; out_features |] dtype in\n let bias = bias_init.f [| out_features |] dtype in\n {\n params =\n Ptree.dict\n [ (\"weight\", Ptree.tensor weight); (\"bias\", Ptree.tensor bias) ];\n state = Ptree.empty;\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.linear\" dtype x in\n let fields = Ptree.Dict.fields_exn ~ctx:\"Layer.linear\" params in\n let weight = Ptree.Dict.get_tensor_exn fields ~name:\"weight\" dtype in\n let bias = Ptree.Dict.get_tensor_exn fields ~name:\"bias\" dtype in\n (Nx.add (Nx.matmul x weight) bias, state));\n }\n\n(* Convolution *)\n\nlet conv1d ~in_channels ~out_channels ?(kernel_size = 3) ?(stride = 1)\n ?(dilation = 1) ?(padding = `Same) () =\n let weight_init = Init.glorot_uniform ~in_axis:1 ~out_axis:0 () in\n {\n init =\n (fun ~dtype ->\n let weight =\n weight_init.f [| out_channels; in_channels; kernel_size |] dtype\n in\n let bias = Nx.zeros dtype [| out_channels |] in\n {\n params =\n Ptree.dict\n [ (\"weight\", Ptree.tensor weight); (\"bias\", Ptree.tensor bias) ];\n state = Ptree.empty;\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.conv1d\" dtype x in\n let fields = Ptree.Dict.fields_exn ~ctx:\"Layer.conv1d\" params in\n let weight = Ptree.Dict.get_tensor_exn fields ~name:\"weight\" dtype in\n let bias = Ptree.Dict.get_tensor_exn fields ~name:\"bias\" dtype in\n let x =\n match padding with\n | `Same | `Valid -> x\n | `Causal ->\n let pad_left = (kernel_size - 1) * dilation in\n Nx.pad [| (0, 0); (0, 0); (pad_left, 0) |] 0.0 x\n in\n let padding =\n match padding with `Same -> `Same | `Valid | `Causal -> `Valid\n in\n (Fn.conv1d ~stride ~dilation ~padding ~bias x weight, state));\n }\n\nlet conv2d ~in_channels ~out_channels ?(kernel_size = (3, 3)) () =\n let kh, kw = kernel_size in\n let weight_init = Init.glorot_uniform ~in_axis:1 ~out_axis:0 () in\n {\n init =\n (fun ~dtype ->\n let weight =\n weight_init.f [| out_channels; in_channels; kh; kw |] dtype\n in\n let bias = Nx.zeros dtype [| out_channels |] in\n {\n params =\n Ptree.dict\n [ (\"weight\", Ptree.tensor weight); (\"bias\", Ptree.tensor bias) ];\n state = Ptree.empty;\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.conv2d\" dtype x in\n let fields = Ptree.Dict.fields_exn ~ctx:\"Layer.conv2d\" params in\n let weight = Ptree.Dict.get_tensor_exn fields ~name:\"weight\" dtype in\n let bias = Ptree.Dict.get_tensor_exn fields ~name:\"bias\" dtype in\n (Fn.conv2d ~padding:`Same ~bias x weight, state));\n }\n\n(* Normalization *)\n\nlet layer_norm ~dim ?(eps = 1e-5) () =\n {\n init =\n (fun ~dtype ->\n let gamma = Nx.ones dtype [| dim |] in\n let beta = Nx.zeros dtype [| dim |] in\n {\n params =\n Ptree.dict\n [ (\"gamma\", Ptree.tensor gamma); (\"beta\", Ptree.tensor beta) ];\n state = Ptree.empty;\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.layer_norm\" dtype x in\n let fields = Ptree.Dict.fields_exn ~ctx:\"Layer.layer_norm\" params in\n let gamma = Ptree.Dict.get_tensor_exn fields ~name:\"gamma\" dtype in\n let beta = Ptree.Dict.get_tensor_exn fields ~name:\"beta\" dtype in\n (Fn.layer_norm ~gamma ~beta ~epsilon:eps x, state));\n }\n\nlet rms_norm ~dim ?(eps = 1e-6) () =\n {\n init =\n (fun ~dtype ->\n let scale = Nx.ones dtype [| dim |] in\n {\n params = Ptree.dict [ (\"scale\", Ptree.tensor scale) ];\n state = Ptree.empty;\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.rms_norm\" dtype x in\n let fields = Ptree.Dict.fields_exn ~ctx:\"Layer.rms_norm\" params in\n let scale = Ptree.Dict.get_tensor_exn fields ~name:\"scale\" dtype in\n (Fn.rms_norm ~gamma:scale ~epsilon:eps x, state));\n }\n\nlet batch_norm ~num_features () =\n {\n init =\n (fun ~dtype ->\n let scale = Nx.ones dtype [| num_features |] in\n let bias = Nx.zeros dtype [| num_features |] in\n let running_mean = Nx.zeros dtype [| num_features |] in\n let running_var = Nx.ones dtype [| num_features |] in\n {\n params =\n Ptree.dict\n [ (\"scale\", Ptree.tensor scale); (\"bias\", Ptree.tensor bias) ];\n state =\n Ptree.dict\n [\n (\"running_mean\", Ptree.tensor running_mean);\n (\"running_var\", Ptree.tensor running_var);\n ];\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore ctx;\n let x = require_same_float_dtype ~ctx:\"Layer.batch_norm\" dtype x in\n let params_fields =\n Ptree.Dict.fields_exn ~ctx:\"Layer.batch_norm.params\" params\n in\n let state_fields =\n Ptree.Dict.fields_exn ~ctx:\"Layer.batch_norm.state\" state\n in\n let scale =\n Ptree.Dict.get_tensor_exn params_fields ~name:\"scale\" dtype\n in\n let bias = Ptree.Dict.get_tensor_exn params_fields ~name:\"bias\" dtype in\n let running_mean =\n Ptree.Dict.get_tensor_exn state_fields ~name:\"running_mean\" dtype\n in\n let running_var =\n Ptree.Dict.get_tensor_exn state_fields ~name:\"running_var\" dtype\n in\n let rank = Array.length (Nx.shape x) in\n let axes =\n match rank with\n | 2 -> [ 0 ]\n | 3 -> [ 0; 2 ]\n | 4 -> [ 0; 2; 3 ]\n | _ -> [ 0 ]\n in\n if training then\n let momentum = 0.99 in\n let one_minus = 1.0 -. momentum in\n let batch_mean = Nx.mean ~axes x in\n let batch_var = Nx.var ~axes x in\n let y = Fn.batch_norm ~axes ~scale ~bias x in\n let running_mean' =\n Nx.add\n (Nx.mul running_mean (Nx.scalar dtype momentum))\n (Nx.mul batch_mean (Nx.scalar dtype one_minus))\n in\n let running_var' =\n Nx.add\n (Nx.mul running_var (Nx.scalar dtype momentum))\n (Nx.mul batch_var (Nx.scalar dtype one_minus))\n in\n let state' =\n Ptree.dict\n [\n (\"running_mean\", Ptree.tensor running_mean');\n (\"running_var\", Ptree.tensor running_var');\n ]\n in\n (y, state')\n else\n let scale_eval, bias_eval =\n match rank with\n | 2 ->\n ( Nx.reshape [| 1; num_features |] scale,\n Nx.reshape [| 1; num_features |] bias )\n | 3 ->\n ( Nx.reshape [| 1; num_features; 1 |] scale,\n Nx.reshape [| 1; num_features; 1 |] bias )\n | 4 ->\n ( Nx.reshape [| 1; num_features; 1; 1 |] scale,\n Nx.reshape [| 1; num_features; 1; 1 |] bias )\n | _ -> (scale, bias)\n in\n let y =\n Nx.standardize ~axes ~mean:running_mean ~variance:running_var x\n |> fun normalized -> Nx.add (Nx.mul normalized scale_eval) bias_eval\n in\n (y, state));\n }\n\n(* Embedding *)\n\nlet embedding ~vocab_size ~embed_dim ?(scale = true) () =\n let emb_init = Init.normal ~stddev:0.02 () in\n {\n init =\n (fun ~dtype ->\n let embedding = emb_init.f [| vocab_size; embed_dim |] dtype in\n {\n params = Ptree.dict [ (\"embedding\", Ptree.tensor embedding) ];\n state = Ptree.empty;\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx indices ->\n ignore (training, ctx);\n let fields = Ptree.Dict.fields_exn ~ctx:\"Layer.embedding\" params in\n let embedding =\n Ptree.Dict.get_tensor_exn fields ~name:\"embedding\" dtype\n in\n let indices = require_int32_indices ~ctx:\"Layer.embedding\" indices in\n (Fn.embedding ~scale ~embedding indices, state));\n }\n\n(* Regularization *)\n\nlet dropout ~rate () =\n if rate < 0.0 || rate >= 1.0 then\n invalid_argf \"Layer.dropout: expected 0.0 <= rate < 1.0, got %g\" rate;\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.dropout\" dtype x in\n if (not training) || rate = 0.0 then (x, state)\n else (Fn.dropout ~rate x, state));\n }\n\n(* Activation layers *)\n\nlet relu () =\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.relu\" dtype x in\n (Nx.relu x, state));\n }\n\nlet gelu () =\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.gelu\" dtype x in\n (Activation.gelu x, state));\n }\n\nlet silu () =\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.silu\" dtype x in\n (Activation.silu x, state));\n }\n\nlet tanh () =\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.tanh\" dtype x in\n (Nx.tanh x, state));\n }\n\nlet sigmoid () =\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.sigmoid\" dtype x in\n (Nx.sigmoid x, state));\n }\n\n(* Pooling *)\n\nlet max_pool2d ~kernel_size ?stride () =\n let stride = match stride with Some value -> value | None -> kernel_size in\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.max_pool2d\" dtype x in\n (Fn.max_pool2d ~kernel_size ~stride x, state));\n }\n\nlet avg_pool2d ~kernel_size ?stride () =\n let stride = match stride with Some value -> value | None -> kernel_size in\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.avg_pool2d\" dtype x in\n (Fn.avg_pool2d ~kernel_size ~stride x, state));\n }\n\n(* Reshape *)\n\nlet flatten () =\n {\n init = (fun ~dtype -> { params = Ptree.empty; state = Ptree.empty; dtype });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx x ->\n ignore (params, training, ctx);\n let x = require_same_float_dtype ~ctx:\"Layer.flatten\" dtype x in\n let shape = Nx.shape x in\n if Array.length shape = 0 then\n invalid_arg \"Layer.flatten: expected rank >= 1\";\n let batch_size = shape.(0) in\n let flat_size =\n Array.fold_left ( * ) 1 (Array.sub shape 1 (Array.length shape - 1))\n in\n let x = if Nx.is_c_contiguous x then x else Nx.contiguous x in\n (Nx.reshape [| batch_size; flat_size |] x, state));\n }\n\n(* Composition *)\n\nlet sequential layers =\n {\n init =\n (fun ~dtype ->\n let n = List.length layers in\n let keys = Array.init n (fun _ -> Nx.Rng.next_key ()) in\n let acc_params = ref [] in\n let acc_state = ref [] in\n List.iteri\n (fun i module_ ->\n let vars =\n Nx.Rng.with_key keys.(i) (fun () -> module_.init ~dtype)\n in\n acc_params := vars.params :: !acc_params;\n acc_state := vars.state :: !acc_state)\n layers;\n {\n params = Ptree.list (List.rev !acc_params);\n state = Ptree.list (List.rev !acc_state);\n dtype;\n });\n apply =\n (fun ~params ~state ~dtype ~training ?ctx input ->\n let params_items =\n Ptree.List.items_exn ~ctx:\"Layer.sequential.params\" params\n in\n let state_items =\n Ptree.List.items_exn ~ctx:\"Layer.sequential.state\" state\n in\n match (layers, params_items, state_items) with\n | [], [], [] ->\n let input =\n require_same_float_dtype ~ctx:\"Layer.sequential\" dtype input\n in\n (input, Ptree.list [])\n | first :: rest_layers, p :: ps, s :: ss ->\n let y, first_state =\n first.apply ~params:p ~state:s ~dtype ~training ?ctx input\n in\n let rec go modules param_values state_values x =\n match (modules, param_values, state_values) with\n | [], [], [] -> (x, [])\n | module_ :: modules_tail, p :: ps_tail, s :: ss_tail ->\n let y, state' =\n module_.apply ~params:p ~state:s ~dtype ~training ?ctx x\n in\n let y_final, state_tail = go modules_tail ps_tail ss_tail y in\n (y_final, state' :: state_tail)\n | _ ->\n invalid_arg\n \"Layer.sequential: params/state/layers length mismatch\"\n in\n let y_final, rest_states = go rest_layers ps ss y in\n (y_final, Ptree.list (first_state :: rest_states))\n | _ ->\n invalid_arg \"Layer.sequential: params/state/layers length mismatch\");\n }\n"
},
{
"path": "packages/kaun/lib/layer.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Composable neural network layers.\n\n A {!type:t} pairs parameter/state initialization with a forward computation.\n Layers compose with {!compose} for heterogeneous pipelines (for example\n embeddings to dense layers) and with {!sequential} for homogeneous float\n pipelines. *)\n\n(** {1:types Types} *)\n\ntype 'layout vars\n(** The type for model variables.\n\n [params] are trainable variables consumed by {!Optim}. [state] is\n non-trainable mutable state updated by forward passes (for example running\n statistics in {!batch_norm}). *)\n\nval params : 'layout vars -> Ptree.t\n(** [params v] is [v]'s trainable parameter tree. *)\n\nval state : 'layout vars -> Ptree.t\n(** [state v] is [v]'s non-trainable mutable state tree. *)\n\nval dtype : 'layout vars -> (float, 'layout) Nx.dtype\n(** [dtype v] is [v]'s floating dtype witness. *)\n\nval with_params : 'layout vars -> Ptree.t -> 'layout vars\n(** [with_params v params] is [v] with replaced trainable parameters. *)\n\nval with_state : 'layout vars -> Ptree.t -> 'layout vars\n(** [with_state v state] is [v] with replaced non-trainable state. *)\n\nval make_vars :\n params:Ptree.t ->\n state:Ptree.t ->\n dtype:(float, 'layout) Nx.dtype ->\n 'layout vars\n(** [make_vars ~params ~state ~dtype] builds model variables.\n\n This is mainly useful for layer constructors implemented outside the\n {!Layer} module. *)\n\ntype ('input, 'output) t = {\n init : 'layout. dtype:(float, 'layout) Nx.dtype -> 'layout vars;\n apply :\n 'layout 'in_elt.\n params:Ptree.t ->\n state:Ptree.t ->\n dtype:(float, 'layout) Nx.dtype ->\n training:bool ->\n ?ctx:Context.t ->\n ('input, 'in_elt) Nx.t ->\n ('output, 'layout) Nx.t * Ptree.t;\n}\n(** The type for layers.\n\n [init] creates fresh [params] and [state]. [apply] computes a forward pass\n and returns updated [state]. Random operations (weight initialization,\n dropout) use the implicit RNG scope established by {!Nx.Rng.run} or\n {!Nx.Rng.with_key}.\n\n The input tensor's dtype witness ['in_elt] is independent of the model's\n float dtype witness ['layout]. This allows layers like {!embedding} to\n accept [int32_elt] indices while the model parameters use [float32_elt].\n Float-consuming layers (e.g. {!linear}) require the input dtype to match the\n model dtype exactly and raise [Invalid_argument] on mismatch.\n\n [ctx] carries per-call auxiliary data (attention masks, position ids,\n encoder memory). Most layers ignore it; transformer layers read from it\n using well-known key names. See {!Context}. *)\n\nval init : ('a, 'b) t -> dtype:(float, 'layout) Nx.dtype -> 'layout vars\n(** [init m ~dtype] is [m]'s fresh variables.\n\n Composite layers ({!compose}, {!sequential}) isolate sub-network RNG streams\n via {!Nx.Rng.with_key}. *)\n\nval apply :\n ('a, 'b) t ->\n 'layout vars ->\n training:bool ->\n ?ctx:Context.t ->\n ('a, 'in_elt) Nx.t ->\n ('b, 'layout) Nx.t * 'layout vars\n(** [apply m vars ~training ?ctx x] is the forward pass of [m].\n\n Returns [(y, vars')] where [params vars' = params vars] and [state vars'] is\n the updated state from the forward pass.\n\n The input tensor's dtype witness ['in_elt] is independent of the model's\n float dtype witness ['layout]. For float-consuming layers, the input must\n have the same dtype as the model; a mismatch raises [Invalid_argument].\n\n [training] controls stochastic/stateful behavior. For example, {!dropout}\n uses dropout masks only when [training = true], and {!batch_norm} updates\n running statistics only when [training = true].\n\n [ctx] carries per-call auxiliary data such as attention masks. See\n {!Context}. *)\n\n(** {1:compose Composition} *)\n\nval compose : ('a, 'b) t -> ('b, 'c) t -> ('a, 'c) t\n(** [compose left right] applies [left] then [right].\n\n Parameters and state are stored as {!Ptree.Dict} nodes with keys [\"left\"]\n and [\"right\"]. The RNG key is split between both layers during\n initialization and forward pass. *)\n\nval sequential : (float, float) t list -> (float, float) t\n(** [sequential layers] applies [layers] in order.\n\n Parameters and state are stored as {!Ptree.List} nodes with one entry per\n layer. The RNG key is split per layer during initialization and forward\n pass.\n\n Raises [Invalid_argument] if runtime parameter/state list lengths do not\n match [layers]. *)\n\n(** {1:dense Dense} *)\n\nval linear :\n in_features:int ->\n out_features:int ->\n ?weight_init:Init.t ->\n ?bias_init:Init.t ->\n unit ->\n (float, float) t\n(** [linear ~in_features ~out_features ?weight_init ?bias_init ()] is the fully\n connected map [xW + b].\n\n [weight_init] defaults to {!Init.glorot_uniform ()}. [bias_init] defaults to\n {!Init.zeros}.\n\n Parameters:\n - [weight] with shape [[in_features; out_features]].\n - [bias] with shape [[out_features]]. *)\n\n(** {1:conv Convolution} *)\n\nval conv1d :\n in_channels:int ->\n out_channels:int ->\n ?kernel_size:int ->\n ?stride:int ->\n ?dilation:int ->\n ?padding:[ `Same | `Valid | `Causal ] ->\n unit ->\n (float, float) t\n(** [conv1d ~in_channels ~out_channels ?kernel_size ?stride ?dilation ?padding\n ()] is 1D convolution over inputs shaped [[batch; in_channels; length]].\n\n [kernel_size] defaults to [3]. [stride] defaults to [1]. [dilation] defaults\n to [1]. [padding] defaults to [`Same].\n\n Parameters:\n - [weight] with shape [[out_channels; in_channels; kernel_size]].\n - [bias] with shape [[out_channels]]. *)\n\nval conv2d :\n in_channels:int ->\n out_channels:int ->\n ?kernel_size:int * int ->\n unit ->\n (float, float) t\n(** [conv2d ~in_channels ~out_channels ?kernel_size ()] is 2D convolution over\n inputs shaped [[batch; in_channels; height; width]].\n\n [kernel_size] defaults to [(3, 3)]. Stride is fixed to [(1, 1)] and padding\n mode is [`Same].\n\n Parameters:\n - [weight] with shape [[out_channels; in_channels; kh; kw]].\n - [bias] with shape [[out_channels]]. *)\n\n(** {1:norm Normalization} *)\n\nval layer_norm : dim:int -> ?eps:float -> unit -> (float, float) t\n(** [layer_norm ~dim ?eps ()] is layer normalization with learnable affine\n parameters.\n\n [eps] defaults to [1e-5].\n\n Parameters:\n - [gamma] with shape [[dim]].\n - [beta] with shape [[dim]]. *)\n\nval rms_norm : dim:int -> ?eps:float -> unit -> (float, float) t\n(** [rms_norm ~dim ?eps ()] is RMS normalization with learnable scale.\n\n [eps] defaults to [1e-6].\n\n Parameters:\n - [scale] with shape [[dim]]. *)\n\nval batch_norm : num_features:int -> unit -> (float, float) t\n(** [batch_norm ~num_features ()] is stateful batch normalization.\n\n During training, batch statistics are used and running statistics are\n updated. During evaluation, running statistics are used and preserved.\n\n Normalization axes are inferred from rank:\n - rank 2 uses [[0]].\n - rank 3 uses [[0; 2]].\n - rank 4 uses [[0; 2; 3]].\n - other ranks use [[0]].\n\n Parameters:\n - [scale] with shape [[num_features]].\n - [bias] with shape [[num_features]].\n\n State:\n - [running_mean] with shape [[num_features]].\n - [running_var] with shape [[num_features]]. *)\n\n(** {1:embed Embedding} *)\n\nval embedding :\n vocab_size:int -> embed_dim:int -> ?scale:bool -> unit -> (int32, float) t\n(** [embedding ~vocab_size ~embed_dim ?scale ()] is an embedding lookup layer.\n\n Inputs are int32 token indices. Output shape is\n [indices_shape ++ [embed_dim]].\n\n [scale] defaults to [true]. When [true], output vectors are multiplied by\n [sqrt embed_dim].\n\n Parameters:\n - [embedding] with shape [[vocab_size; embed_dim]]. *)\n\n(** {1:reg Regularization} *)\n\nval dropout : rate:float -> unit -> (float, float) t\n(** [dropout ~rate ()] is elementwise dropout.\n\n When [training = false], it is identity. When [training = true], dropout\n masks are generated using keys from the implicit RNG scope.\n\n Raises [Invalid_argument] if [rate] is outside [0.0 <= rate < 1.0]. *)\n\n(** {1:act Activation Layers} *)\n\nval relu : unit -> (float, float) t\n(** [relu ()] is [max(0, x)]. No parameters. *)\n\nval gelu : unit -> (float, float) t\n(** [gelu ()] is the Gaussian error linear unit. No parameters. *)\n\nval silu : unit -> (float, float) t\n(** [silu ()] is [x * sigmoid(x)]. No parameters. *)\n\nval tanh : unit -> (float, float) t\n(** [tanh ()] is hyperbolic tangent. No parameters. *)\n\nval sigmoid : unit -> (float, float) t\n(** [sigmoid ()] is the logistic function. No parameters. *)\n\n(** {1:pool Pooling} *)\n\nval max_pool2d :\n kernel_size:int * int -> ?stride:int * int -> unit -> (float, float) t\n(** [max_pool2d ~kernel_size ?stride ()] is 2D max pooling.\n\n [stride] defaults to [kernel_size]. No parameters. *)\n\nval avg_pool2d :\n kernel_size:int * int -> ?stride:int * int -> unit -> (float, float) t\n(** [avg_pool2d ~kernel_size ?stride ()] is 2D average pooling.\n\n [stride] defaults to [kernel_size]. No parameters. *)\n\n(** {1:reshape Reshape} *)\n\nval flatten : unit -> (float, float) t\n(** [flatten ()] flattens all dimensions after the batch dimension.\n\n [[batch; d1; ...; dn]] becomes [[batch; d1 * ... * dn]].\n\n Raises [Invalid_argument] if the input rank is [0]. *)\n"
},
{
"path": "packages/kaun/lib/loss.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\n\nlet invalid_argf_fn fn fmt =\n Printf.ksprintf (fun msg -> invalid_argf \"Loss.%s: %s\" fn msg) fmt\n\nlet check_logits_shape ~fn logits =\n let logits_shape = Nx.shape logits in\n let logits_rank = Array.length logits_shape in\n if logits_rank < 1 then invalid_argf_fn fn \"logits must have rank >= 1\";\n let class_axis = logits_rank - 1 in\n let num_classes = logits_shape.(class_axis) in\n if num_classes <= 0 then\n invalid_argf_fn fn \"logits class dimension must be positive (got %d)\"\n num_classes;\n logits_shape\n\nlet check_same_shape ~fn ~rhs_name lhs rhs =\n let lhs_rank = Array.length lhs in\n let rhs_rank = Array.length rhs in\n if rhs_rank <> lhs_rank then\n invalid_argf_fn fn \"%s rank mismatch (got %d, expected %d)\" rhs_name\n rhs_rank lhs_rank;\n for i = 0 to lhs_rank - 1 do\n if rhs.(i) <> lhs.(i) then\n invalid_argf_fn fn \"%s shape mismatch at axis %d (got %d, expected %d)\"\n rhs_name i rhs.(i) lhs.(i)\n done\n\nlet check_cross_entropy_shapes logits labels =\n let fn = \"cross_entropy\" in\n let logits_shape = check_logits_shape ~fn logits in\n let labels_shape = Nx.shape labels in\n check_same_shape ~fn ~rhs_name:\"labels\" logits_shape labels_shape\n\nlet cross_entropy logits labels =\n check_cross_entropy_shapes logits labels;\n let max_logits = Nx.max logits ~axes:[ -1 ] ~keepdims:true in\n let shifted = Nx.sub logits max_logits in\n let log_sum_exp =\n Nx.log (Nx.sum (Nx.exp shifted) ~axes:[ -1 ] ~keepdims:true)\n in\n let log_softmax = Nx.sub shifted log_sum_exp in\n let per_example = Nx.neg (Nx.sum (Nx.mul labels log_softmax) ~axes:[ -1 ]) in\n Nx.mean per_example\n\nlet check_sparse_indices_dtype indices =\n let fn = \"cross_entropy_sparse\" in\n let dtype = Nx.dtype indices in\n if not (Nx_core.Dtype.is_int dtype) then\n invalid_argf_fn fn \"expected integer labels, got %s\"\n (Nx_core.Dtype.to_string dtype)\n\nlet check_sparse_shapes logits indices =\n let fn = \"cross_entropy_sparse\" in\n let logits_shape = check_logits_shape ~fn logits in\n let indices_shape = Nx.shape indices in\n let logits_rank = Array.length logits_shape in\n let indices_rank = Array.length indices_shape in\n if indices_rank <> logits_rank - 1 then\n invalid_argf_fn fn \"labels rank mismatch (got %d, expected %d)\" indices_rank\n (logits_rank - 1);\n for i = 0 to indices_rank - 1 do\n if indices_shape.(i) <> logits_shape.(i) then\n invalid_argf_fn fn\n \"labels shape mismatch at axis %d (got %d, expected %d)\" i\n indices_shape.(i) logits_shape.(i)\n done;\n let class_axis = logits_rank - 1 in\n logits_shape.(class_axis)\n\nlet cross_entropy_sparse logits indices =\n check_sparse_indices_dtype indices;\n ignore (check_sparse_shapes logits indices : int);\n let indices_int = Nx.cast Nx.int32 indices in\n (* Numerically stable log-softmax *)\n let max_logits = Nx.max logits ~axes:[ -1 ] ~keepdims:true in\n let shifted = Nx.sub logits max_logits in\n let log_sum_exp =\n Nx.log (Nx.sum (Nx.exp shifted) ~axes:[ -1 ] ~keepdims:true)\n in\n (* Gather true-class logits: [...] → [...; 1] for take_along_axis *)\n let indices_expanded = Nx.expand_dims [ -1 ] indices_int in\n let true_logits = Nx.take_along_axis ~axis:(-1) indices_expanded shifted in\n (* loss = -(true_logit - log_sum_exp) *)\n let per_example =\n Nx.neg\n (Nx.sub\n (Nx.squeeze ~axes:[ -1 ] true_logits)\n (Nx.squeeze ~axes:[ -1 ] log_sum_exp))\n in\n Nx.mean per_example\n\nlet binary_cross_entropy logits labels =\n let fn = \"binary_cross_entropy\" in\n let logits_shape = Nx.shape logits in\n let labels_shape = Nx.shape labels in\n check_same_shape ~fn ~rhs_name:\"labels\" logits_shape labels_shape;\n let dtype = Nx.dtype logits in\n let one = Nx.scalar dtype 1.0 in\n let log_p = Activation.log_sigmoid logits in\n let log_1_minus_p = Activation.log_sigmoid (Nx.neg logits) in\n let per_element =\n Nx.neg\n (Nx.add (Nx.mul labels log_p) (Nx.mul (Nx.sub one labels) log_1_minus_p))\n in\n Nx.mean per_element\n\nlet mse predictions targets =\n let diff = Nx.sub predictions targets in\n Nx.mean (Nx.mul diff diff)\n\nlet mae predictions targets = Nx.mean (Nx.abs (Nx.sub predictions targets))\n"
},
{
"path": "packages/kaun/lib/loss.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Loss functions.\n\n Losses are differentiable through Rune's autodiff and return scalar means.\n [Invalid_argument] messages are prefixed with [Loss.:]. *)\n\n(** {1:classification Classification} *)\n\nval cross_entropy : (float, 'a) Nx.t -> (float, 'a) Nx.t -> (float, 'a) Nx.t\n(** [cross_entropy logits one_hot_labels] is softmax cross-entropy.\n\n [logits] has shape [[...; num_classes]] and must be rank >= 1.\n [one_hot_labels] must have the same shape.\n\n Uses the log-sum-exp trick for numerical stability.\n\n Raises [Invalid_argument] if ranks or shapes differ, or if [num_classes] is\n not positive. *)\n\nval cross_entropy_sparse : (float, 'a) Nx.t -> ('c, 'd) Nx.t -> (float, 'a) Nx.t\n(** [cross_entropy_sparse logits class_indices] is {!cross_entropy} with integer\n labels.\n\n [class_indices] has shape [[...]] and must match [logits] without the last\n dimension. The class dimension is [logits]' last axis.\n\n Raises [Invalid_argument] if labels are non-integer, ranks mismatch,\n non-class dimensions differ, or the class dimension is non-positive. *)\n\nval binary_cross_entropy :\n (float, 'a) Nx.t -> (float, 'a) Nx.t -> (float, 'a) Nx.t\n(** [binary_cross_entropy logits labels] is sigmoid binary cross-entropy.\n\n [logits] are raw (not sigmoid-normalized). [labels] are typically in\n \\[[0];[1]\\]. Uses log-sigmoid for numerical stability.\n\n Raises [Invalid_argument] if [logits] and [labels] shapes differ. *)\n\n(** {1:regression Regression} *)\n\nval mse : ('a, 'b) Nx.t -> ('a, 'b) Nx.t -> ('a, 'b) Nx.t\n(** [mse predictions targets] is [mean ((predictions - targets)^2)].\n\n Shape compatibility follows Nx broadcasting semantics. *)\n\nval mae : ('a, 'b) Nx.t -> ('a, 'b) Nx.t -> ('a, 'b) Nx.t\n(** [mae predictions targets] is [mean (abs (predictions - targets))].\n\n Shape compatibility follows Nx broadcasting semantics. *)\n"
},
{
"path": "packages/kaun/lib/metric.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet strf = Printf.sprintf\n\n(* Tracker *)\n\ntype entry = { mutable sum : float; mutable n : int }\ntype tracker = (string, entry) Hashtbl.t\n\nlet tracker () : tracker = Hashtbl.create 16\n\nlet observe (t : tracker) name value =\n match Hashtbl.find_opt t name with\n | Some e ->\n e.sum <- e.sum +. value;\n e.n <- e.n + 1\n | None -> Hashtbl.replace t name { sum = value; n = 1 }\n\nlet find_entry t name = Hashtbl.find t name\n\nlet mean t name =\n let e = find_entry t name in\n e.sum /. float_of_int e.n\n\nlet count t name =\n let e = find_entry t name in\n e.n\n\nlet reset t = Hashtbl.reset t\n\nlet to_list t =\n let pairs =\n Hashtbl.fold (fun k e acc -> (k, e.sum /. float_of_int e.n) :: acc) t []\n in\n List.sort (fun (a, _) (b, _) -> String.compare a b) pairs\n\nlet summary t =\n let pairs = to_list t in\n String.concat \" \" (List.map (fun (k, v) -> strf \"%s: %.4f\" k v) pairs)\n\n(* Dataset evaluation *)\n\nlet eval f data =\n let sum = ref 0.0 in\n let n = ref 0 in\n Data.iter\n (fun x ->\n sum := !sum +. f x;\n incr n)\n data;\n if !n = 0 then invalid_arg \"Metric.eval: empty dataset\";\n !sum /. float_of_int !n\n\nlet eval_many f data =\n let tbl = Hashtbl.create 8 in\n let n = ref 0 in\n Data.iter\n (fun x ->\n let pairs = f x in\n List.iter\n (fun (k, v) ->\n match Hashtbl.find_opt tbl k with\n | Some e ->\n e.sum <- e.sum +. v;\n e.n <- e.n + 1\n | None -> Hashtbl.replace tbl k { sum = v; n = 1 })\n pairs;\n incr n)\n data;\n if !n = 0 then invalid_arg \"Metric.eval_many: empty dataset\";\n let pairs =\n Hashtbl.fold (fun k e acc -> (k, e.sum /. float_of_int e.n) :: acc) tbl []\n in\n List.sort (fun (a, _) (b, _) -> String.compare a b) pairs\n\ntype average = Macro | Micro | Weighted\n\n(* Metric functions *)\n\nlet accuracy (type a b c) (predictions : (float, a) Nx.t)\n (targets : (b, c) Nx.t) =\n let pred_shape = Nx.shape predictions in\n let rank = Array.length pred_shape in\n let predicted =\n if rank >= 2 then\n (* Multi-class: argmax along last axis *)\n Nx.argmax ~axis:(-1) predictions\n else\n (* Binary: threshold at 0.5 *)\n let half = Nx.scalar (Nx.dtype predictions) 0.5 in\n Nx.cast Nx.int32 (Nx.greater predictions half)\n in\n let targets_i32 = Nx.cast Nx.int32 targets in\n let correct = Nx.equal predicted targets_i32 in\n let correct_f = Nx.cast Nx.float32 correct in\n Nx.item [] (Nx.mean correct_f)\n\nlet binary_accuracy ?(threshold = 0.5) predictions targets =\n let dtype = Nx.dtype predictions in\n let thresh = Nx.scalar dtype threshold in\n let predicted = Nx.cast Nx.float32 (Nx.greater predictions thresh) in\n let targets_f = Nx.cast Nx.float32 targets in\n let correct = Nx.equal predicted targets_f in\n let correct_f = Nx.cast Nx.float32 correct in\n Nx.item [] (Nx.mean correct_f)\n\n(* Classification metrics *)\n\nlet confusion_counts (type a b c) (predictions : (float, a) Nx.t)\n (targets : (b, c) Nx.t) =\n let pred_shape = Nx.shape predictions in\n let num_classes = pred_shape.(Array.length pred_shape - 1) in\n let predicted = Nx.argmax ~axis:(-1) predictions in\n let targets_i32 = Nx.cast Nx.int32 targets in\n let pred_oh = Nx.cast Nx.float32 (Nx.one_hot ~num_classes predicted) in\n let tgt_oh = Nx.cast Nx.float32 (Nx.one_hot ~num_classes targets_i32) in\n let tp = Nx.sum (Nx.mul pred_oh tgt_oh) ~axes:[ 0 ] in\n let pred_sum = Nx.sum pred_oh ~axes:[ 0 ] in\n let tgt_sum = Nx.sum tgt_oh ~axes:[ 0 ] in\n let fp = Nx.sub pred_sum tp in\n let fn = Nx.sub tgt_sum tp in\n (tp, fp, fn, num_classes)\n\nlet safe_div a b = if b = 0.0 then 0.0 else a /. b\n\nlet precision avg predictions targets =\n let tp, fp, fn, num_classes = confusion_counts predictions targets in\n let tp = Nx.to_array tp in\n let fp = Nx.to_array fp in\n match avg with\n | Micro ->\n let tp_sum = Array.fold_left ( +. ) 0.0 tp in\n let fp_sum = Array.fold_left ( +. ) 0.0 fp in\n safe_div tp_sum (tp_sum +. fp_sum)\n | Macro ->\n let sum = ref 0.0 in\n for c = 0 to num_classes - 1 do\n sum := !sum +. safe_div tp.(c) (tp.(c) +. fp.(c))\n done;\n !sum /. float_of_int num_classes\n | Weighted ->\n let fn = Nx.to_array fn in\n let w_sum = ref 0.0 in\n let total = ref 0.0 in\n for c = 0 to num_classes - 1 do\n let support = tp.(c) +. fn.(c) in\n w_sum := !w_sum +. (support *. safe_div tp.(c) (tp.(c) +. fp.(c)));\n total := !total +. support\n done;\n safe_div !w_sum !total\n\nlet recall avg predictions targets =\n let tp, _fp, fn, num_classes = confusion_counts predictions targets in\n let tp = Nx.to_array tp in\n let fn = Nx.to_array fn in\n match avg with\n | Micro ->\n let tp_sum = Array.fold_left ( +. ) 0.0 tp in\n let fn_sum = Array.fold_left ( +. ) 0.0 fn in\n safe_div tp_sum (tp_sum +. fn_sum)\n | Macro ->\n let sum = ref 0.0 in\n for c = 0 to num_classes - 1 do\n sum := !sum +. safe_div tp.(c) (tp.(c) +. fn.(c))\n done;\n !sum /. float_of_int num_classes\n | Weighted ->\n let w_sum = ref 0.0 in\n let total = ref 0.0 in\n for c = 0 to num_classes - 1 do\n let support = tp.(c) +. fn.(c) in\n w_sum := !w_sum +. (support *. safe_div tp.(c) (tp.(c) +. fn.(c)));\n total := !total +. support\n done;\n safe_div !w_sum !total\n\nlet f1 avg predictions targets =\n let tp, fp, fn, num_classes = confusion_counts predictions targets in\n let tp = Nx.to_array tp in\n let fp = Nx.to_array fp in\n let fn = Nx.to_array fn in\n match avg with\n | Micro ->\n let tp_sum = Array.fold_left ( +. ) 0.0 tp in\n let fp_sum = Array.fold_left ( +. ) 0.0 fp in\n let fn_sum = Array.fold_left ( +. ) 0.0 fn in\n safe_div (2.0 *. tp_sum) ((2.0 *. tp_sum) +. fp_sum +. fn_sum)\n | Macro ->\n let sum = ref 0.0 in\n for c = 0 to num_classes - 1 do\n sum :=\n !sum +. safe_div (2.0 *. tp.(c)) ((2.0 *. tp.(c)) +. fp.(c) +. fn.(c))\n done;\n !sum /. float_of_int num_classes\n | Weighted ->\n let w_sum = ref 0.0 in\n let total = ref 0.0 in\n for c = 0 to num_classes - 1 do\n let support = tp.(c) +. fn.(c) in\n w_sum :=\n !w_sum\n +. support\n *. safe_div (2.0 *. tp.(c)) ((2.0 *. tp.(c)) +. fp.(c) +. fn.(c));\n total := !total +. support\n done;\n safe_div !w_sum !total\n"
},
{
"path": "packages/kaun/lib/metric.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Training metrics.\n\n {!Metric} provides running scalar tracking and dataset evaluation.\n\n A {!type:tracker} accumulates named running means during training. For\n dataset evaluation, {!eval} and {!eval_many} fold user-supplied functions\n over a {!Data.t} pipeline and return averaged results.\n\n Metric functions such as {!accuracy} are plain tensor-to-scalar functions\n that compose freely with {!eval}. *)\n\n(** {1:tracker Running Tracker} *)\n\ntype tracker\n(** A mutable set of named running-mean accumulators. *)\n\nval tracker : unit -> tracker\n(** [tracker ()] is a fresh tracker with no observations. *)\n\nval observe : tracker -> string -> float -> unit\n(** [observe t name value] records [value] under [name]. *)\n\nval mean : tracker -> string -> float\n(** [mean t name] is the running mean of observations under [name].\n\n Raises [Not_found] if [name] was never observed. *)\n\nval count : tracker -> string -> int\n(** [count t name] is the number of observations under [name].\n\n Raises [Not_found] if [name] was never observed. *)\n\nval reset : tracker -> unit\n(** [reset t] clears all observations. *)\n\nval to_list : tracker -> (string * float) list\n(** [to_list t] is the current means as [(name, mean)] pairs, sorted by name. *)\n\nval summary : tracker -> string\n(** [summary t] is a human-readable one-liner of all current means, e.g.\n [\"accuracy: 0.9150 loss: 0.4231\"]. *)\n\n(** {1:eval Dataset Evaluation} *)\n\nval eval : ('a -> float) -> 'a Data.t -> float\n(** [eval f data] is the mean of [f batch] over all elements of [data].\n\n Raises [Invalid_argument] if [data] yields no elements. *)\n\nval eval_many :\n ('a -> (string * float) list) -> 'a Data.t -> (string * float) list\n(** [eval_many f data] is the per-name mean of [f batch] over all elements of\n [data]. Returns [(name, mean)] pairs sorted by name.\n\n Raises [Invalid_argument] if [data] yields no elements. *)\n\n(** {1:average Averaging} *)\n\ntype average =\n | Macro\n | Micro\n | Weighted\n (** The type for multi-class averaging modes.\n - [Macro] is the unweighted mean of per-class scores.\n - [Micro] aggregates TP, FP, FN globally before computing.\n - [Weighted] is the mean of per-class scores weighted by class support\n (number of true instances). *)\n\n(** {1:compute Common Metric Functions} *)\n\nval accuracy : (float, 'a) Nx.t -> ('b, 'c) Nx.t -> float\n(** [accuracy predictions targets] is the fraction of correct predictions.\n\n Multi-class: [predictions] has shape [[batch; num_classes]] (logits or\n probabilities), [targets] has shape [[batch]] (integer class indices).\n Predicted class is [argmax] along the last axis.\n\n Binary: both tensors have shape [[batch]] or [[batch; 1]]. Predictions above\n [0.5] count as class [1]. *)\n\nval binary_accuracy :\n ?threshold:float -> (float, 'a) Nx.t -> (float, 'a) Nx.t -> float\n(** [binary_accuracy ?threshold predictions targets] is the fraction of correct\n binary predictions.\n\n [threshold] defaults to [0.5]. Predictions above [threshold] count as class\n [1], targets are expected in \\[[0];[1]\\]. *)\n\n(** {1:classification Classification} *)\n\nval precision : average -> (float, 'a) Nx.t -> ('b, 'c) Nx.t -> float\n(** [precision avg predictions targets] is the precision score.\n\n [predictions] has shape [[batch; num_classes]] (logits or probabilities).\n [targets] has shape [[batch]] (integer class indices). Predicted class is\n [argmax] along the last axis.\n\n When a class has no predicted instances, its precision is [0.0]. *)\n\nval recall : average -> (float, 'a) Nx.t -> ('b, 'c) Nx.t -> float\n(** [recall avg predictions targets] is the recall score.\n\n Input convention is the same as {!precision}.\n\n When a class has no true instances, its recall is [0.0]. *)\n\nval f1 : average -> (float, 'a) Nx.t -> ('b, 'c) Nx.t -> float\n(** [f1 avg predictions targets] is the F1 score (harmonic mean of {!precision}\n and {!recall}).\n\n Input convention is the same as {!precision}.\n\n When both precision and recall are [0.0] for a class, its F1 is [0.0]. *)\n"
},
{
"path": "packages/kaun/lib/optim.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule Dtype = Nx_core.Dtype\n\n(* Helpers *)\n\nlet err_expected_float_dtype = \"Optim: expected floating-point dtype\"\n\nlet float_of_scalar (type a b) (dtype : (a, b) Dtype.t) (value : a) : float =\n match dtype with\n | Dtype.Float16 ->\n let value : float = value in\n value\n | Dtype.Float32 ->\n let value : float = value in\n value\n | Dtype.Float64 ->\n let value : float = value in\n value\n | Dtype.BFloat16 ->\n let value : float = value in\n value\n | Dtype.Float8_e4m3 ->\n let value : float = value in\n value\n | Dtype.Float8_e5m2 ->\n let value : float = value in\n value\n | _ -> invalid_arg err_expected_float_dtype\n\nlet scalar dt x = Nx.scalar dt (Dtype.of_float dt x)\n\nlet tensor_sum_sq (Ptree.P t) =\n let dtype = Nx.dtype t in\n let sq = Nx.mul t t in\n float_of_scalar dtype (Nx.item [] (Nx.sum sq))\n\n(* Per-leaf packed Vega state with captured dtype for type unification *)\n\ntype packed_vega_state =\n | PVS : {\n dtype : ('a, 'b) Dtype.t;\n st : ('a, 'b) Vega.state;\n }\n -> packed_vega_state\n\n(* State *)\n\ntype state = { tx : Vega.t; leaf_states : packed_vega_state array }\n\n(* Init *)\n\nlet init tx params =\n let leaves, _ = Ptree.flatten params in\n let leaf_states =\n Array.of_list\n (List.map\n (fun pt ->\n Ptree.with_tensor pt\n {\n run = (fun t -> PVS { dtype = Nx.dtype t; st = Vega.init tx t });\n })\n leaves)\n in\n { tx; leaf_states }\n\n(* Update: returns updates tree (not new params) *)\n\nlet update st params grads =\n let param_leaves, rebuild = Ptree.flatten params in\n let grad_leaves, _ = Ptree.flatten grads in\n let n = Array.length st.leaf_states in\n let update_packed = Array.make n (List.hd param_leaves) in\n let new_leaf_states = Array.make n st.leaf_states.(0) in\n List.iteri\n (fun i param_pt ->\n let grad_pt = List.nth grad_leaves i in\n let (PVS { dtype = dt; st = vega_st }) = st.leaf_states.(i) in\n let param_t = Ptree.Tensor.to_typed_exn dt param_pt in\n let grad_t = Ptree.Tensor.to_typed_exn dt grad_pt in\n let upd, new_vega_st = Vega.update vega_st ~grad:grad_t ~param:param_t in\n update_packed.(i) <- Ptree.P upd;\n new_leaf_states.(i) <- PVS { dtype = dt; st = new_vega_st })\n param_leaves;\n let updates = rebuild (Array.to_list update_packed) in\n (updates, { tx = st.tx; leaf_states = new_leaf_states })\n\n(* Apply updates: add updates to params *)\n\nlet apply_updates params updates =\n Ptree.map2 { run = (fun param upd -> Nx.add param upd) } params updates\n\n(* Step: convenience for update + apply_updates *)\n\nlet step st params grads =\n let updates, new_st = update st params grads in\n let new_params = apply_updates params updates in\n (new_params, new_st)\n\n(* Serialization *)\n\nlet state_to_trees st =\n let n = Array.length st.leaf_states in\n if n = 0 then (0, [])\n else\n (* Get count from first leaf (all leaves share the same count) *)\n let (PVS { st = first_st; _ }) = st.leaf_states.(0) in\n let count, _ = Vega.state_to_tensors first_st in\n (* Extract per-leaf tensor arrays *)\n let per_leaf_tensors = Array.make n [||] in\n for i = 0 to n - 1 do\n let (PVS { st = vega_st; _ }) = st.leaf_states.(i) in\n let _, tensors = Vega.state_to_tensors vega_st in\n per_leaf_tensors.(i) <- Array.map (fun t -> Ptree.P t) tensors\n done;\n (* Determine number of state tensors per leaf *)\n let n_tensors = Array.length per_leaf_tensors.(0) in\n if n_tensors = 0 then (count, [])\n else\n (* Transpose: per-leaf x per-tensor -> per-tensor x per-leaf *)\n let tensor_trees =\n List.init n_tensors (fun m ->\n let leaves =\n List.init n (fun i -> Ptree.Tensor per_leaf_tensors.(i).(m))\n in\n Ptree.List leaves)\n in\n (count, tensor_trees)\n\nlet state_of_trees tx ~count trees =\n let n_trees = List.length trees in\n let expected_tensors = Vega.n_tensors tx in\n if n_trees <> expected_tensors then\n invalid_arg\n (Printf.sprintf \"Optim.state_of_trees: expected %d moment trees, got %d\"\n expected_tensors n_trees);\n if n_trees = 0 then { tx; leaf_states = [||] }\n else\n let first_tree = List.hd trees in\n let first_items = Ptree.List.items_exn first_tree in\n let n_leaves = List.length first_items in\n (* Collect per-tensor leaf lists *)\n let tensor_leaves =\n List.map\n (fun tree -> List.map Ptree.as_tensor_exn (Ptree.List.items_exn tree))\n trees\n in\n (* Transpose: per-tensor x per-leaf -> per-leaf x per-tensor *)\n let leaf_states =\n Array.init n_leaves (fun i ->\n let leaf_tensors =\n List.map (fun moment -> List.nth moment i) tensor_leaves\n in\n (* Use the first tensor's dtype as reference *)\n let ref_pt = List.hd leaf_tensors in\n Ptree.with_tensor ref_pt\n {\n run =\n (fun ref_t ->\n let dt = Nx.dtype ref_t in\n let typed_tensors =\n Array.of_list\n (List.map (Ptree.Tensor.to_typed_exn dt) leaf_tensors)\n in\n let vega_st = Vega.state_of_tensors tx ~count typed_tensors in\n PVS { dtype = dt; st = vega_st });\n })\n in\n { tx; leaf_states }\n\n(* Gradient utilities *)\n\nlet global_norm t =\n let sum_sq = Ptree.fold (fun acc p -> acc +. tensor_sum_sq p) 0. t in\n sqrt sum_sq\n\nlet clip_by_global_norm max_norm grads =\n let norm = global_norm grads in\n if norm <= max_norm then grads\n else\n let scale = max_norm /. norm in\n Ptree.map\n {\n run =\n (fun t ->\n let dt = Nx.dtype t in\n Nx.mul t (scalar dt scale));\n }\n grads\n"
},
{
"path": "packages/kaun/lib/optim.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Parameter-tree optimizer.\n\n Bridges {!Vega} with {!Ptree}: applies per-parameter optimizer steps across\n all leaves of a heterogeneous parameter tree. *)\n\n(** {1:types Types} *)\n\ntype state\n(** Optimizer state for a parameter tree. Packs per-leaf {!Vega} states. *)\n\n(** {1:core Core} *)\n\nval init : Vega.t -> Ptree.t -> state\n(** [init tx params] initializes optimizer state for all leaves of [params]. *)\n\nval update : state -> Ptree.t -> Ptree.t -> Ptree.t * state\n(** [update state params grads] returns [(updates, new_state)].\n\n Applies {!Vega.update} to each matching leaf pair. The returned [updates]\n tree has the same structure as [params] and can be applied via\n {!apply_updates}. *)\n\nval apply_updates : Ptree.t -> Ptree.t -> Ptree.t\n(** [apply_updates params updates] adds [updates] to [params] element-wise\n across all leaves. *)\n\nval step : state -> Ptree.t -> Ptree.t -> Ptree.t * state\n(** [step state params grads] returns [(new_params, new_state)].\n\n Convenience for:\n {[\n let updates, state = update state params grads in\n (apply_updates params updates, state)\n ]} *)\n\n(** {1:serialization Serialization} *)\n\nval state_to_trees : state -> int * Ptree.t list\n(** [state_to_trees st] is [(count, trees)] where [count] is the optimizer step\n count and [trees] are the internal state as parameter trees.\n\n Transforms with no state tensors return an empty list. *)\n\nval state_of_trees : Vega.t -> count:int -> Ptree.t list -> state\n(** [state_of_trees tx ~count trees] reconstructs optimizer state from a\n transformation, step count, and serialized trees.\n\n Raises [Invalid_argument] if the number of trees does not match the\n transformation's expectation. *)\n\n(** {1:grad Gradient Utilities} *)\n\nval clip_by_global_norm : float -> Ptree.t -> Ptree.t\n(** [clip_by_global_norm max_norm grads] rescales [grads] so their global L2\n norm does not exceed [max_norm]. Returns [grads] unchanged if the norm is\n already within bounds.\n\n Raises [Invalid_argument] if a leaf tensor is not floating point. *)\n\nval global_norm : Ptree.t -> float\n(** [global_norm t] is the L2 norm across all leaf tensors of [t].\n\n Raises [Invalid_argument] if a leaf tensor is not floating point. *)\n"
},
{
"path": "packages/kaun/lib/ptree.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype tensor = P : ('a, 'layout) Nx.t -> tensor\ntype t = Tensor of tensor | List of t list | Dict of (string * t) list\n\nlet invalid_argf fmt = Printf.ksprintf invalid_arg fmt\n\nlet invalid_arg_ctx ?ctx msg =\n match ctx with\n | None -> invalid_arg msg\n | Some ctx -> invalid_argf \"%s: %s\" ctx msg\n\nlet expected ?ctx what = invalid_arg_ctx ?ctx (\"expected \" ^ what)\n\nlet key_not_found ?ctx key =\n match ctx with\n | None -> invalid_argf \"key %S not found\" key\n | Some ctx -> invalid_argf \"%s: key %S not found\" ctx key\n\nlet tensor x = Tensor (P x)\nlet list xs = List xs\nlet empty = List []\n\nlet validate_key key =\n if String.length key = 0 then invalid_arg \"empty key\";\n for i = 0 to String.length key - 1 do\n let c = String.unsafe_get key i in\n match c with\n | '.' | '[' | ']' ->\n invalid_argf\n \"key %S contains reserved character %C (keys must not contain '.', \\\n '[', or ']')\"\n key c\n | _ -> ()\n done\n\nlet dict kvs =\n let tbl = Hashtbl.create (Stdlib.List.length kvs) in\n Stdlib.List.iter\n (fun (k, _) ->\n validate_key k;\n if Hashtbl.mem tbl k then invalid_argf \"duplicate key %S\" k\n else Hashtbl.add tbl k ())\n kvs;\n Dict kvs\n\nmodule Tensor = struct\n let dtype (P t) = Nx_core.Dtype.pack (Nx.dtype t)\n let shape (P t) = Nx.shape t\n let numel (P t) = Nx.numel t\n\n let to_typed (type a l) (dtype : (a, l) Nx.dtype) (P t) : (a, l) Nx.t option =\n match Nx_core.Dtype.equal_witness (Nx.dtype t) dtype with\n | Some Type.Equal -> Some t\n | None -> None\n\n let to_typed_exn (type a l) (dtype : (a, l) Nx.dtype) (P t) : (a, l) Nx.t =\n match Nx_core.Dtype.equal_witness (Nx.dtype t) dtype with\n | Some Type.Equal -> t\n | None ->\n invalid_argf \"dtype mismatch: expected %s, got %s\"\n (Nx_core.Dtype.to_string dtype)\n (Nx_core.Dtype.to_string (Nx.dtype t))\nend\n\nmodule Dict = struct\n type fields = (string * t) list\n\n let fields_exn ?ctx t =\n match t with Dict kvs -> kvs | _ -> expected ?ctx \"Dict\"\n\n let find key fields = Stdlib.List.assoc_opt key fields\n\n let find_exn ?ctx key fields =\n match find key fields with Some v -> v | None -> key_not_found ?ctx key\n\n let get_tensor_exn fields ~name dtype =\n match find_exn name fields with\n | Tensor p -> Tensor.to_typed_exn dtype p\n | _ -> invalid_argf \"field %S is not a tensor\" name\nend\n\nmodule List = struct\n let items_exn ?ctx t =\n match t with List xs -> xs | _ -> expected ?ctx \"List\"\nend\n\ntype 'r tensor_handler = { run : 'a 'layout. ('a, 'layout) Nx.t -> 'r }\n\ntype map_handler = {\n run : 'a 'layout. ('a, 'layout) Nx.t -> ('a, 'layout) Nx.t;\n}\n\ntype map2_handler = {\n run :\n 'a 'layout. ('a, 'layout) Nx.t -> ('a, 'layout) Nx.t -> ('a, 'layout) Nx.t;\n}\n\nlet with_tensor (P t) (handler : _ tensor_handler) = handler.run t\n\nlet as_tensor_exn ?ctx t =\n match t with Tensor p -> p | _ -> expected ?ctx \"Tensor\"\n\nlet map (f : map_handler) t =\n let rec go = function\n | Tensor (P x) -> Tensor (P (f.run x))\n | List xs -> List (Stdlib.List.map go xs)\n | Dict kvs -> Dict (Stdlib.List.map (fun (k, v) -> (k, go v)) kvs)\n in\n go t\n\nlet map2 (f : map2_handler) a b =\n let rec go a b =\n match (a, b) with\n | Tensor (P x), Tensor (P y) -> (\n match Nx_core.Dtype.equal_witness (Nx.dtype x) (Nx.dtype y) with\n | Some Type.Equal -> Tensor (P (f.run x y))\n | None -> invalid_arg \"dtype mismatch\")\n | List xs, List ys ->\n if Stdlib.List.length xs <> Stdlib.List.length ys then\n invalid_arg \"list length mismatch\";\n List (Stdlib.List.map2 go xs ys)\n | Dict kvs1, Dict kvs2 ->\n if Stdlib.List.length kvs1 <> Stdlib.List.length kvs2 then\n invalid_arg \"dict size mismatch\";\n Dict\n (Stdlib.List.map\n (fun (k, v1) ->\n match Stdlib.List.assoc_opt k kvs2 with\n | Some v2 -> (k, go v1 v2)\n | None -> invalid_argf \"key %S not found in second dict\" k)\n kvs1)\n | _ -> invalid_arg \"structure mismatch\"\n in\n go a b\n\nlet iter f t =\n let rec go = function\n | Tensor p -> f p\n | List xs -> Stdlib.List.iter go xs\n | Dict kvs -> Stdlib.List.iter (fun (_, v) -> go v) kvs\n in\n go t\n\nlet fold f acc t =\n let rec go acc = function\n | Tensor p -> f acc p\n | List xs -> Stdlib.List.fold_left go acc xs\n | Dict kvs -> Stdlib.List.fold_left (fun acc (_, v) -> go acc v) acc kvs\n in\n go acc t\n\nlet flatten t =\n let tensors = ref [] in\n iter (fun p -> tensors := p :: !tensors) t;\n let tensors = Stdlib.List.rev !tensors in\n let rebuild new_tensors =\n let remaining = ref new_tensors in\n let take () =\n match !remaining with\n | [] -> invalid_arg \"not enough tensors to rebuild tree\"\n | x :: rest ->\n remaining := rest;\n x\n in\n let rec go = function\n | Tensor _ -> Tensor (take ())\n | List xs -> List (Stdlib.List.map go xs)\n | Dict kvs -> Dict (Stdlib.List.map (fun (k, v) -> (k, go v)) kvs)\n in\n let result = go t in\n (match !remaining with\n | [] -> ()\n | _ -> invalid_arg \"too many tensors to rebuild tree\");\n result\n in\n (tensors, rebuild)\n\nlet flatten_with_paths t =\n let join prefix seg = if prefix = \"\" then seg else prefix ^ \".\" ^ seg in\n let acc = ref [] in\n let rec go prefix = function\n | Tensor p -> acc := (prefix, p) :: !acc\n | List xs ->\n Stdlib.List.iteri (fun i v -> go (join prefix (string_of_int i)) v) xs\n | Dict kvs -> Stdlib.List.iter (fun (k, v) -> go (join prefix k) v) kvs\n in\n go \"\" t;\n Stdlib.List.rev !acc\n\nlet zeros_like t = map { run = Nx.zeros_like } t\nlet count_parameters t = fold (fun acc p -> acc + Tensor.numel p) 0 t\n\nlet pp_shape shape =\n Stdlib.String.concat \"x\"\n (Stdlib.Array.to_list (Stdlib.Array.map string_of_int shape))\n\nlet rec pp_with_indent indent ppf = function\n | Tensor p ->\n with_tensor p\n {\n run =\n (fun t ->\n Format.fprintf ppf \"Tensor(%s, %s)\"\n (Nx_core.Dtype.to_string (Nx.dtype t))\n (pp_shape (Nx.shape t)));\n }\n | List [] -> Format.pp_print_string ppf \"List []\"\n | List xs ->\n let next_indent = indent ^ \" \" in\n Format.pp_print_string ppf \"List [\";\n Stdlib.List.iter\n (fun v ->\n Format.pp_print_char ppf '\\n';\n Format.pp_print_string ppf next_indent;\n pp_with_indent next_indent ppf v)\n xs;\n Format.pp_print_char ppf '\\n';\n Format.pp_print_string ppf indent;\n Format.pp_print_char ppf ']'\n | Dict [] -> Format.pp_print_string ppf \"Dict {}\"\n | Dict kvs ->\n let next_indent = indent ^ \" \" in\n Format.pp_print_string ppf \"Dict {\";\n Stdlib.List.iter\n (fun (k, v) ->\n Format.pp_print_char ppf '\\n';\n Format.pp_print_string ppf next_indent;\n Format.pp_print_string ppf k;\n Format.pp_print_string ppf \": \";\n pp_with_indent next_indent ppf v)\n kvs;\n Format.pp_print_char ppf '\\n';\n Format.pp_print_string ppf indent;\n Format.pp_print_char ppf '}'\n\nlet pp ppf t = pp_with_indent \"\" ppf t\n"
},
{
"path": "packages/kaun/lib/ptree.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Heterogeneous parameter trees.\n\n A parameter tree is a finite tree with tensor leaves and container nodes.\n Leaves are packed tensors ({!tensor}), and containers are either ordered\n lists ([List]) or string-keyed dicts ([Dict]). *)\n\ntype tensor =\n | P : ('a, 'layout) Nx.t -> tensor\n (** A packed tensor. The wrapper hides dtype and layout parameters. *)\n\ntype t =\n | Tensor of tensor (** A tensor leaf. *)\n | List of t list (** An ordered list branch. *)\n | Dict of (string * t) list (** A dict branch. Keys are strings. *)\n\n(** {1:constructors Constructors} *)\n\nval tensor : ('a, 'layout) Nx.t -> t\n(** [tensor x] is [Tensor (P x)]. *)\n\nval list : t list -> t\n(** [list xs] is [List xs]. *)\n\nval dict : (string * t) list -> t\n(** [dict kvs] is [Dict kvs] with key validation.\n\n Raises [Invalid_argument] if a key is empty, duplicated, or contains ['.'],\n ['\\['], or ['\\]']. *)\n\nval empty : t\n(** [empty] is [List []]. Canonical value for \"no parameters\" or \"no state\". *)\n\n(** {1:tensor Tensor Inspection} *)\n\nmodule Tensor : sig\n val dtype : tensor -> Nx_core.Dtype.packed\n (** [dtype t] is [t]'s dtype. *)\n\n val shape : tensor -> int array\n (** [shape t] is [t]'s shape. *)\n\n val numel : tensor -> int\n (** [numel t] is the number of elements in [t]. *)\n\n val to_typed : ('a, 'l) Nx.dtype -> tensor -> ('a, 'l) Nx.t option\n (** [to_typed dtype t] is [Some x] iff [t] has dtype [dtype], with [x] the\n typed tensor. It is [None] on dtype mismatch. *)\n\n val to_typed_exn : ('a, 'l) Nx.dtype -> tensor -> ('a, 'l) Nx.t\n (** [to_typed_exn dtype t] is the typed tensor in [t].\n\n Raises [Invalid_argument] if [t]'s dtype is not [dtype]. *)\nend\n\n(** {1:dict Dict Access} *)\n\nmodule Dict : sig\n type fields = (string * t) list\n (** The type for dict fields. *)\n\n val fields_exn : ?ctx:string -> t -> fields\n (** [fields_exn ?ctx t] is [t]'s fields.\n\n Raises [Invalid_argument] if [t] is not [Dict _]. The optional [ctx] is\n prefixed to the error message. *)\n\n val find : string -> fields -> t option\n (** [find name fields] is [Some v] if [name] is bound in [fields], and [None]\n otherwise. *)\n\n val find_exn : ?ctx:string -> string -> fields -> t\n (** [find_exn ?ctx name fields] is [name]'s value in [fields].\n\n Raises [Invalid_argument] if [name] is missing. The optional [ctx] is\n prefixed to the error message. *)\n\n val get_tensor_exn :\n fields -> name:string -> ('a, 'l) Nx_core.Dtype.t -> ('a, 'l) Nx.t\n (** [get_tensor_exn fields ~name dtype] is the typed tensor in [fields] under\n [name].\n\n Raises [Invalid_argument] if [name] is missing, [name] is not a tensor, or\n the tensor dtype differs from [dtype]. *)\nend\n\n(** {1:list List Access} *)\n\nmodule List : sig\n val items_exn : ?ctx:string -> t -> t list\n (** [items_exn ?ctx t] is [t]'s items.\n\n Raises [Invalid_argument] if [t] is not [List _]. The optional [ctx] is\n prefixed to the error message. *)\nend\n\n(** {1:leaf Leaf Access} *)\n\ntype 'r tensor_handler = { run : 'a 'layout. ('a, 'layout) Nx.t -> 'r }\n(** Rank-2 handler for unpacking {!tensor}. *)\n\nval with_tensor : tensor -> 'a tensor_handler -> 'a\n(** [with_tensor t h] applies [h.run] to the unpacked tensor in [t]. *)\n\nval as_tensor_exn : ?ctx:string -> t -> tensor\n(** [as_tensor_exn ?ctx t] is [t]'s packed tensor.\n\n Raises [Invalid_argument] if [t] is not [Tensor _]. The optional [ctx] is\n prefixed to the error message. *)\n\n(** {1:functional Functional Operations} *)\n\ntype map_handler = {\n run : 'a 'layout. ('a, 'layout) Nx.t -> ('a, 'layout) Nx.t;\n}\n(** Rank-2 tensor mapper used by {!map}. *)\n\nval map : map_handler -> t -> t\n(** [map f t] maps [f.run] over tensor leaves and preserves tree structure. *)\n\ntype map2_handler = {\n run :\n 'a 'layout. ('a, 'layout) Nx.t -> ('a, 'layout) Nx.t -> ('a, 'layout) Nx.t;\n}\n(** Rank-2 tensor zipper used by {!map2}. *)\n\nval map2 : map2_handler -> t -> t -> t\n(** [map2 f a b] zips [a] and [b] and applies [f.run] to paired tensor leaves.\n\n Lists are matched by position. Dict nodes are matched by key using [a]'s key\n order.\n\n Raises [Invalid_argument] on structure mismatch, list or dict size mismatch,\n missing keys in [b], or paired dtype mismatch. *)\n\nval iter : (tensor -> unit) -> t -> unit\n(** [iter f t] applies [f] to each leaf tensor in depth-first order.\n\n Leaves are visited left-to-right in list order and dict field order. *)\n\nval fold : ('acc -> tensor -> 'acc) -> 'acc -> t -> 'acc\n(** [fold f acc t] folds leaf tensors in the same traversal order as {!iter}. *)\n\n(** {1:flatten Flatten and Rebuild} *)\n\nval flatten : t -> tensor list * (tensor list -> t)\n(** [flatten t] is [(leaves, rebuild)] where:\n - [leaves] are [t]'s leaf tensors in depth-first order;\n - [rebuild new_leaves] rebuilds [t]'s structure with [new_leaves].\n\n [rebuild] raises [Invalid_argument] if [new_leaves] has a different length\n than [leaves]. *)\n\nval flatten_with_paths : t -> (string * tensor) list\n(** [flatten_with_paths t] returns [(path, tensor)] pairs where paths are\n dot-separated strings. Dict keys become path segments; list indices become\n decimal segments (e.g. [\"layers.0.weight\"]).\n\n If [t] is a tensor leaf, its path is the empty string.\n\n The path encoding is injective for trees built with {!dict}, because {!dict}\n rejects keys containing ['.'], ['\\['], or ['\\]']. *)\n\n(** {1:utils Utilities} *)\n\nval zeros_like : t -> t\n(** [zeros_like t] has the same structure as [t], with each tensor replaced by\n [Nx.zeros_like]. *)\n\nval count_parameters : t -> int\n(** [count_parameters t] is the sum of {!Tensor.numel} over all leaf tensors. *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp] formats trees for debugging. *)\n"
},
{
"path": "packages/kaun/lib/train.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype ('i, 'o) t = { model : ('i, 'o) Layer.t; optimizer : Vega.t }\ntype 'l state = { vars : 'l Layer.vars; opt_state : Optim.state }\n\nlet make ~model ~optimizer = { model; optimizer }\n\nlet init t ~dtype =\n let vars = Layer.init t.model ~dtype in\n let opt_state = Optim.init t.optimizer (Layer.params vars) in\n { vars; opt_state }\n\nlet vars st = st.vars\n\nlet make_state t vars =\n let opt_state = Optim.init t.optimizer (Layer.params vars) in\n { vars; opt_state }\n\nlet step (type i o l in_elt) (t : (i, o) t) (st : l state) ~training ?ctx\n ~(loss : (o, l) Nx.t -> (float, l) Nx.t) (x : (i, in_elt) Nx.t) =\n let loss_val, grads, new_layer_state =\n Grad.value_and_grad_aux\n (fun params ->\n let vars' = Layer.with_params st.vars params in\n let pred, vars'' = Layer.apply t.model vars' ~training ?ctx x in\n (loss pred, Layer.state vars''))\n (Layer.params st.vars)\n in\n let new_params, opt_state =\n Optim.step st.opt_state (Layer.params st.vars) grads\n in\n let vars =\n Layer.with_params st.vars new_params |> fun v ->\n Layer.with_state v new_layer_state\n in\n (loss_val, { vars; opt_state })\n\nexception Early_stop\n\nlet fit (type i o l in_elt) (t : (i, o) t) (st : l state) ?ctx ?report\n (data : ((i, in_elt) Nx.t * ((o, l) Nx.t -> (float, l) Nx.t)) Data.t) =\n let st = ref st in\n let i = ref 0 in\n (try\n Data.iter\n (fun (x, loss) ->\n incr i;\n let loss_val, st' = step t !st ~training:true ?ctx ~loss x in\n st := st';\n match report with\n | Some f -> f ~step:!i ~loss:(Nx.item [] loss_val) !st\n | None -> ())\n data\n with Early_stop -> ());\n !st\n\nlet predict (type i o l in_elt) (t : (i, o) t) (st : l state) ?ctx\n (x : (i, in_elt) Nx.t) =\n let y, _ = Layer.apply t.model st.vars ~training:false ?ctx x in\n y\n"
},
{
"path": "packages/kaun/lib/train.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** High-level training loop.\n\n {!Train} composes {!Layer}, {!Grad}, and {!Optim} into a single training\n driver. Users never touch parameter trees, optimizer state, or gradient\n computation directly.\n\n For advanced use, {!step} exposes a single training step and {!vars} gives\n access to the underlying model variables. *)\n\n(** {1:types Types} *)\n\ntype ('i, 'o) t\n(** The type for trainers. A trainer pairs a model with an optimizer. *)\n\ntype 'l state\n(** The type for training state. Bundles model variables and optimizer state. *)\n\n(** {1:core Core} *)\n\nval make : model:('i, 'o) Layer.t -> optimizer:Vega.t -> ('i, 'o) t\n(** [make ~model ~optimizer] creates a trainer. *)\n\nval init : ('i, 'o) t -> dtype:(float, 'l) Nx.dtype -> 'l state\n(** [init trainer ~dtype] initializes model variables and optimizer state.\n\n Random keys for weight initialization are drawn from the implicit RNG scope.\n*)\n\nval vars : 'l state -> 'l Layer.vars\n(** [vars st] is the current model variables (params + state + dtype). *)\n\nval make_state : ('i, 'o) t -> 'l Layer.vars -> 'l state\n(** [make_state trainer vars] is a training state with [vars] and freshly\n initialized optimizer state.\n\n Use this to start training from pretrained or externally loaded weights\n instead of {!init}. *)\n\n(** {1:training Training} *)\n\nexception Early_stop\n(** Raise inside [report] to end training early. {!fit} catches this exception\n and returns the current state. *)\n\nval step :\n ('i, 'o) t ->\n 'l state ->\n training:bool ->\n ?ctx:Context.t ->\n loss:(('o, 'l) Nx.t -> (float, 'l) Nx.t) ->\n ('i, 'in_elt) Nx.t ->\n (float, 'l) Nx.t * 'l state\n(** [step trainer st ~training ?ctx ~loss x] performs one training step.\n\n Computes the forward pass, differentiates the loss with respect to trainable\n parameters, applies the optimizer, and threads updated layer state.\n\n [ctx] is forwarded to the model's forward pass. See {!Context}.\n\n When [training = false], gradients are still computed and optimizer is still\n applied. Use {!predict} for pure inference. *)\n\nval fit :\n ('i, 'o) t ->\n 'l state ->\n ?ctx:Context.t ->\n ?report:(step:int -> loss:float -> 'l state -> unit) ->\n (('i, 'in_elt) Nx.t * (('o, 'l) Nx.t -> (float, 'l) Nx.t)) Data.t ->\n 'l state\n(** [fit trainer st ?ctx ?report data] trains the model over [data] and returns\n the final state.\n\n Each element of [data] is a pair [(x, loss_fn)] where [x] is the input\n tensor and [loss_fn] computes the scalar loss from the model output. This\n allows the loss to depend on per-batch labels.\n\n [ctx] is forwarded to the model's forward pass on each step. See {!Context}.\n\n When provided, [report] is called after every step with the step number\n (1-based), scalar loss, and training state. Raise {!Early_stop} inside\n [report] to end training early.\n\n For fixed-data training (same input every step), use {!Data.repeat}:\n {[\n Train.fit trainer st (Data.repeat 1000 (x, loss_fn))\n ]} *)\n\n(** {1:inference Inference} *)\n\nval predict :\n ('i, 'o) t ->\n 'l state ->\n ?ctx:Context.t ->\n ('i, 'in_elt) Nx.t ->\n ('o, 'l) Nx.t\n(** [predict trainer st ?ctx x] runs the model in evaluation mode (no state\n updates, no dropout).\n\n [ctx] is forwarded to the model's forward pass. See {!Context}. *)\n"
},
{
"path": "packages/kaun/test/dune",
"content": "(tests\n (names\n test_ptree\n test_init\n test_loss\n test_layer\n test_fn\n test_optim\n test_grad\n test_attention\n test_data\n test_train\n test_metric\n test_checkpoint)\n (package kaun)\n (libraries rune vega nx nx.core nx.io kaun windtrap))\n"
},
{
"path": "packages/kaun/test/test_attention.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Attention = Kaun.Attention\nmodule Layer = Kaun.Layer\nmodule Ptree = Kaun.Ptree\n\nlet dtype = Nx.float32\n\n(* Init *)\n\nlet test_init_param_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Attention.multi_head_attention ~embed_dim:64 ~num_heads:4 () in\n let vars = Layer.init m ~dtype in\n let fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let shape name =\n Array.to_list (Nx.shape (Ptree.Dict.get_tensor_exn fields ~name dtype))\n in\n equal ~msg:\"q_proj shape\" (list int) [ 64; 64 ] (shape \"q_proj\");\n equal ~msg:\"k_proj shape\" (list int) [ 64; 64 ] (shape \"k_proj\");\n equal ~msg:\"v_proj shape\" (list int) [ 64; 64 ] (shape \"v_proj\");\n equal ~msg:\"out_proj shape\" (list int) [ 64; 64 ] (shape \"out_proj\")\n\nlet test_init_gqa_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m =\n Attention.multi_head_attention ~embed_dim:64 ~num_heads:8 ~num_kv_heads:2 ()\n in\n let vars = Layer.init m ~dtype in\n let fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let shape name =\n Array.to_list (Nx.shape (Ptree.Dict.get_tensor_exn fields ~name dtype))\n in\n let head_dim = 64 / 8 in\n equal ~msg:\"q_proj shape\" (list int) [ 64; 8 * head_dim ] (shape \"q_proj\");\n equal ~msg:\"k_proj shape\" (list int) [ 64; 2 * head_dim ] (shape \"k_proj\");\n equal ~msg:\"v_proj shape\" (list int) [ 64; 2 * head_dim ] (shape \"v_proj\");\n equal ~msg:\"out_proj shape\" (list int) [ 64; 64 ] (shape \"out_proj\")\n\n(* Forward *)\n\nlet test_forward_shape () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Attention.multi_head_attention ~embed_dim:64 ~num_heads:4 () in\n let vars = Layer.init m ~dtype in\n let x = Nx.randn dtype [| 2; 8; 64 |] in\n let y, _vars' = Layer.apply m vars ~training:false x in\n equal ~msg:\"output shape\" (list int) [ 2; 8; 64 ] (Array.to_list (Nx.shape y))\n\nlet test_forward_gqa () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m =\n Attention.multi_head_attention ~embed_dim:64 ~num_heads:8 ~num_kv_heads:2 ()\n in\n let vars = Layer.init m ~dtype in\n let x = Nx.randn dtype [| 2; 8; 64 |] in\n let y, _vars' = Layer.apply m vars ~training:false x in\n equal ~msg:\"GQA output shape\" (list int) [ 2; 8; 64 ]\n (Array.to_list (Nx.shape y))\n\nlet test_causal_differs () =\n Nx.Rng.run ~seed:7 @@ fun () ->\n let m_causal =\n Attention.multi_head_attention ~embed_dim:32 ~num_heads:2 ~is_causal:true ()\n in\n let m_non_causal =\n Attention.multi_head_attention ~embed_dim:32 ~num_heads:2 ~is_causal:false\n ()\n in\n let vars_causal = Layer.init m_causal ~dtype in\n let vars_non_causal = Layer.init m_non_causal ~dtype in\n let x = Nx.randn dtype [| 1; 6; 32 |] in\n let y_causal, _ = Layer.apply m_causal vars_causal ~training:false x in\n let y_non_causal, _ =\n Layer.apply m_non_causal vars_non_causal ~training:false x\n in\n let sum_causal = Nx.item [] (Nx.sum y_causal) in\n let sum_non_causal = Nx.item [] (Nx.sum y_non_causal) in\n is_true ~msg:\"causal vs non-causal differ\"\n (Float.abs (sum_causal -. sum_non_causal) > 1e-6)\n\n(* RoPE *)\n\nlet test_rope_preserves_shape () =\n Nx.Rng.run ~seed:0 @@ fun () ->\n let x = Nx.randn dtype [| 2; 4; 8; 16 |] in\n let y = Attention.rope x in\n equal ~msg:\"rope output shape\" (list int) [ 2; 4; 8; 16 ]\n (Array.to_list (Nx.shape y))\n\nlet test_rope_changes_values () =\n Nx.Rng.run ~seed:0 @@ fun () ->\n let x = Nx.randn dtype [| 1; 2; 4; 8 |] in\n let y = Attention.rope x in\n let diff = Nx.item [] (Nx.sum (Nx.abs (Nx.sub x y))) in\n is_true ~msg:\"rope changes values\" (diff > 0.0)\n\nlet test_rope_seq_dim () =\n Nx.Rng.run ~seed:0 @@ fun () ->\n let x = Nx.randn dtype [| 2; 8; 4; 16 |] in\n let y = Attention.rope ~seq_dim:1 x in\n equal ~msg:\"rope seq_dim shape\" (list int) [ 2; 8; 4; 16 ]\n (Array.to_list (Nx.shape y))\n\nlet test_rope_odd_dim_error () =\n Nx.Rng.run ~seed:0 @@ fun () ->\n let x = Nx.randn dtype [| 1; 2; 4; 7 |] in\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Attention.rope x))\n\n(* Dropout *)\n\nlet test_dropout_eval_identity () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m =\n Attention.multi_head_attention ~embed_dim:32 ~num_heads:2 ~dropout:0.5 ()\n in\n let vars = Layer.init m ~dtype in\n let x = Nx.randn dtype [| 1; 4; 32 |] in\n let y, _ = Layer.apply m vars ~training:false x in\n equal ~msg:\"eval shape\" (list int) [ 1; 4; 32 ] (Array.to_list (Nx.shape y))\n\nlet test_dropout_training () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m =\n Attention.multi_head_attention ~embed_dim:32 ~num_heads:2 ~dropout:0.5 ()\n in\n let vars = Layer.init m ~dtype in\n let x = Nx.randn dtype [| 1; 4; 32 |] in\n let y, _ = Layer.apply m vars ~training:true x in\n equal ~msg:\"training shape\" (list int) [ 1; 4; 32 ]\n (Array.to_list (Nx.shape y))\n\n(* RoPE integration *)\n\nlet test_forward_with_rope () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m =\n Attention.multi_head_attention ~embed_dim:32 ~num_heads:2 ~rope:true ()\n in\n let vars = Layer.init m ~dtype in\n let x = Nx.randn dtype [| 1; 8; 32 |] in\n let y, _ = Layer.apply m vars ~training:false x in\n equal ~msg:\"rope forward shape\" (list int) [ 1; 8; 32 ]\n (Array.to_list (Nx.shape y))\n\nlet () =\n run \"Kaun.Attention\"\n [\n group \"init\"\n [\n test \"param shapes\" test_init_param_shapes;\n test \"GQA param shapes\" test_init_gqa_shapes;\n ];\n group \"forward\"\n [\n test \"output shape\" test_forward_shape;\n test \"GQA output shape\" test_forward_gqa;\n test \"causal differs\" test_causal_differs;\n test \"with RoPE\" test_forward_with_rope;\n ];\n group \"rope\"\n [\n test \"preserves shape\" test_rope_preserves_shape;\n test \"changes values\" test_rope_changes_values;\n test \"respects seq_dim\" test_rope_seq_dim;\n test \"odd dim error\" test_rope_odd_dim_error;\n ];\n group \"dropout\"\n [\n test \"eval identity\" test_dropout_eval_identity;\n test \"training with dropout\" test_dropout_training;\n ];\n ]\n"
},
{
"path": "packages/kaun/test/test_checkpoint.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Checkpoint = Kaun.Checkpoint\nmodule Ptree = Kaun.Ptree\nmodule Optim = Kaun.Optim\n\nlet with_tmpfile f =\n let path = Filename.temp_file \"ckpt\" \".safetensors\" in\n Fun.protect ~finally:(fun () -> Sys.remove path) (fun () -> f path)\n\nlet to_array t = Nx.to_array (Nx.reshape [| -1 |] (Nx.cast Nx.float32 t))\n\n(* Checkpoint save/load *)\n\nlet test_roundtrip_single_tensor () =\n with_tmpfile (fun path ->\n let t = Nx.create Nx.float32 [| 2; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6. |] in\n let tree = Ptree.tensor t in\n Checkpoint.save path tree;\n let loaded = Checkpoint.load path ~like:tree in\n match loaded with\n | Ptree.Tensor (Ptree.P lt) ->\n let vals = to_array lt in\n equal ~msg:\"length\" int 6 (Array.length vals);\n equal ~msg:\"first\" (float 1e-6) 1.0 vals.(0);\n equal ~msg:\"last\" (float 1e-6) 6.0 vals.(5)\n | _ -> fail \"expected Tensor\")\n\nlet test_roundtrip_nested_tree () =\n with_tmpfile (fun path ->\n let w = Nx.create Nx.float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |] in\n let b = Nx.create Nx.float32 [| 2 |] [| 0.1; 0.2 |] in\n let tree =\n Ptree.dict\n [\n ( \"layer0\",\n Ptree.dict\n [ (\"weight\", Ptree.tensor w); (\"bias\", Ptree.tensor b) ] );\n (\"layer1\", Ptree.dict [ (\"weight\", Ptree.tensor w) ]);\n ]\n in\n Checkpoint.save path tree;\n let loaded = Checkpoint.load path ~like:tree in\n let pairs = Ptree.flatten_with_paths loaded in\n equal ~msg:\"num leaves\" int 3 (List.length pairs);\n let names = List.map fst pairs in\n equal ~msg:\"paths\" (list string)\n [ \"layer0.weight\"; \"layer0.bias\"; \"layer1.weight\" ]\n names)\n\nlet test_roundtrip_list_tree () =\n with_tmpfile (fun path ->\n let t0 = Nx.create Nx.float32 [| 3 |] [| 1.; 2.; 3. |] in\n let t1 = Nx.create Nx.float32 [| 2 |] [| 4.; 5. |] in\n let tree = Ptree.list [ Ptree.tensor t0; Ptree.tensor t1 ] in\n Checkpoint.save path tree;\n let loaded = Checkpoint.load path ~like:tree in\n let pairs = Ptree.flatten_with_paths loaded in\n equal ~msg:\"num leaves\" int 2 (List.length pairs);\n let _, Ptree.P lt1 = List.nth pairs 1 in\n let vals = to_array lt1 in\n equal ~msg:\"second tensor\" (float 1e-6) 5.0 vals.(1))\n\nlet test_missing_key () =\n with_tmpfile (fun path ->\n let t = Nx.create Nx.float32 [| 2 |] [| 1.; 2. |] in\n let small = Ptree.dict [ (\"a\", Ptree.tensor t) ] in\n Checkpoint.save path small;\n let big = Ptree.dict [ (\"a\", Ptree.tensor t); (\"b\", Ptree.tensor t) ] in\n raises_invalid_arg \"Checkpoint.load: missing key \\\"b\\\"\" (fun () ->\n ignore (Checkpoint.load path ~like:big)))\n\nlet test_shape_mismatch () =\n with_tmpfile (fun path ->\n let t = Nx.create Nx.float32 [| 2; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6. |] in\n let tree = Ptree.tensor t in\n Checkpoint.save path tree;\n let wrong =\n Ptree.tensor\n (Nx.create Nx.float32 [| 3; 2 |] [| 1.; 2.; 3.; 4.; 5.; 6. |])\n in\n raises_invalid_arg\n \"Checkpoint.load: shape mismatch for \\\"\\\": expected [3; 2], got [2; 3]\"\n (fun () -> ignore (Checkpoint.load path ~like:wrong)))\n\nlet test_dtype_casting () =\n with_tmpfile (fun path ->\n let t = Nx.create Nx.float32 [| 3 |] [| 1.; 2.; 3. |] in\n Checkpoint.save path (Ptree.tensor t);\n let template =\n Ptree.tensor (Nx.create Nx.float64 [| 3 |] [| 0.; 0.; 0. |])\n in\n let loaded = Checkpoint.load path ~like:template in\n match loaded with\n | Ptree.Tensor (Ptree.P lt) ->\n let vals = to_array lt in\n equal ~msg:\"casted value\" (float 1e-6) 2.0 vals.(1)\n | _ -> fail \"expected Tensor\")\n\nlet test_empty_tree () =\n with_tmpfile (fun path ->\n Checkpoint.save path Ptree.empty;\n let loaded = Checkpoint.load path ~like:Ptree.empty in\n match loaded with Ptree.List [] -> () | _ -> fail \"expected empty list\")\n\n(* Optim state serialization *)\n\nlet test_optim_sgd_no_momentum () =\n let params = Ptree.tensor (Nx.create Nx.float32 [| 2 |] [| 1.; 2. |]) in\n let algo = Vega.sgd (Vega.Schedule.constant 0.01) in\n let st = Optim.init algo params in\n let count, trees = Optim.state_to_trees st in\n equal ~msg:\"count\" int 0 count;\n equal ~msg:\"no trees\" int 0 (List.length trees);\n let st' = Optim.state_of_trees algo ~count trees in\n let count', trees' = Optim.state_to_trees st' in\n equal ~msg:\"count roundtrip\" int 0 count';\n equal ~msg:\"trees roundtrip\" int 0 (List.length trees')\n\nlet test_optim_sgd_momentum () =\n let params = Ptree.tensor (Nx.create Nx.float32 [| 2 |] [| 1.; 2. |]) in\n let algo = Vega.sgd ~momentum:0.9 (Vega.Schedule.constant 0.01) in\n let st = Optim.init algo params in\n let count, trees = Optim.state_to_trees st in\n equal ~msg:\"count\" int 0 count;\n equal ~msg:\"one tree\" int 1 (List.length trees)\n\nlet test_optim_adam_roundtrip () =\n let params = Ptree.tensor (Nx.create Nx.float32 [| 2 |] [| 1.; 2. |]) in\n let algo = Vega.adam (Vega.Schedule.constant 0.001) in\n let st = Optim.init algo params in\n let count, trees = Optim.state_to_trees st in\n equal ~msg:\"count\" int 0 count;\n equal ~msg:\"two trees\" int 2 (List.length trees);\n let st' = Optim.state_of_trees algo ~count trees in\n let count', trees' = Optim.state_to_trees st' in\n equal ~msg:\"count roundtrip\" int 0 count';\n equal ~msg:\"trees roundtrip\" int 2 (List.length trees')\n\nlet test_optim_wrong_tree_count () =\n let algo = Vega.adam (Vega.Schedule.constant 0.001) in\n raises_invalid_arg \"Optim.state_of_trees: expected 2 moment trees, got 1\"\n (fun () -> ignore (Optim.state_of_trees algo ~count:0 [ Ptree.empty ]))\n\nlet () =\n run \"Kaun.Checkpoint\"\n [\n group \"save/load\"\n [\n test \"roundtrip single tensor\" test_roundtrip_single_tensor;\n test \"roundtrip nested tree\" test_roundtrip_nested_tree;\n test \"roundtrip list tree\" test_roundtrip_list_tree;\n test \"missing key\" test_missing_key;\n test \"shape mismatch\" test_shape_mismatch;\n test \"dtype casting\" test_dtype_casting;\n test \"empty tree\" test_empty_tree;\n ];\n group \"optim serialization\"\n [\n test \"sgd no momentum\" test_optim_sgd_no_momentum;\n test \"sgd momentum\" test_optim_sgd_momentum;\n test \"adam roundtrip\" test_optim_adam_roundtrip;\n test \"wrong tree count\" test_optim_wrong_tree_count;\n ];\n ]\n"
},
{
"path": "packages/kaun/test/test_data.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Data = Kaun.Data\n\nlet dtype = Nx.float32\n\n(* Constructors *)\n\nlet test_of_array () =\n let d = Data.of_array [| 10; 20; 30 |] in\n equal ~msg:\"length\" (option int) (Some 3) (Data.length d);\n let a = Data.to_array d in\n equal ~msg:\"elements\" (array int) [| 10; 20; 30 |] a\n\nlet test_of_fn () =\n let d = Data.of_fn 4 (fun i -> i * i) in\n equal ~msg:\"length\" (option int) (Some 4) (Data.length d);\n let a = Data.to_array d in\n equal ~msg:\"elements\" (array int) [| 0; 1; 4; 9 |] a\n\nlet test_of_fn_negative () =\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Data.of_fn (-1) Fun.id))\n\nlet test_of_tensor () =\n let t = Nx.create dtype [| 3; 2 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0 |] in\n let d = Data.of_tensor t in\n equal ~msg:\"length\" (option int) (Some 3) (Data.length d);\n let a = Data.to_array d in\n equal ~msg:\"count\" int 3 (Array.length a);\n equal ~msg:\"shape\" (list int) [ 2 ] (Array.to_list (Nx.shape a.(0)));\n equal ~msg:\"first elem\" (float 1e-6) 1.0 (Nx.item [ 0 ] a.(0))\n\nlet test_of_tensors () =\n let x = Nx.create dtype [| 3; 2 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0 |] in\n let y = Nx.create dtype [| 3 |] [| 10.0; 20.0; 30.0 |] in\n let d = Data.of_tensors (x, y) in\n equal ~msg:\"length\" (option int) (Some 3) (Data.length d);\n let a = Data.to_array d in\n equal ~msg:\"count\" int 3 (Array.length a);\n let x0, y0 = a.(0) in\n equal ~msg:\"x0 shape\" (list int) [ 2 ] (Array.to_list (Nx.shape x0));\n equal ~msg:\"y0 scalar\" (float 1e-6) 10.0 (Nx.item [] y0)\n\nlet test_of_tensors_mismatch () =\n let x = Nx.create dtype [| 3; 2 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0 |] in\n let y = Nx.create dtype [| 2 |] [| 10.0; 20.0 |] in\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Data.of_tensors (x, y)))\n\n(* Transformers *)\n\nlet test_map () =\n let d = Data.of_array [| 1; 2; 3 |] |> Data.map (fun x -> x * 2) in\n equal ~msg:\"mapped\" (array int) [| 2; 4; 6 |] (Data.to_array d)\n\nlet test_batch () =\n let d = Data.of_array [| 1; 2; 3; 4; 5 |] |> Data.batch 2 in\n let batches = Data.to_array d in\n equal ~msg:\"num batches\" int 3 (Array.length batches);\n equal ~msg:\"batch 0\" (array int) [| 1; 2 |] batches.(0);\n equal ~msg:\"batch 1\" (array int) [| 3; 4 |] batches.(1);\n equal ~msg:\"batch 2 (partial)\" (array int) [| 5 |] batches.(2)\n\nlet test_batch_drop_last () =\n let d = Data.of_array [| 1; 2; 3; 4; 5 |] |> Data.batch ~drop_last:true 2 in\n let batches = Data.to_array d in\n equal ~msg:\"num batches\" int 2 (Array.length batches);\n equal ~msg:\"batch 0\" (array int) [| 1; 2 |] batches.(0);\n equal ~msg:\"batch 1\" (array int) [| 3; 4 |] batches.(1)\n\nlet test_batch_invalid_size () =\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Data.of_array [| 1; 2 |] |> Data.batch 0))\n\nlet test_map_batch () =\n let d =\n Data.of_array [| 1; 2; 3; 4 |]\n |> Data.map_batch 2 (fun batch -> Array.fold_left ( + ) 0 batch)\n in\n equal ~msg:\"map_batch\" (array int) [| 3; 7 |] (Data.to_array d)\n\nlet test_shuffle_deterministic () =\n let d1 =\n Nx.Rng.run ~seed:42 @@ fun () ->\n Data.of_array [| 0; 1; 2; 3; 4; 5; 6; 7 |] |> Data.shuffle |> Data.to_array\n in\n let d2 =\n Nx.Rng.run ~seed:42 @@ fun () ->\n Data.of_array [| 0; 1; 2; 3; 4; 5; 6; 7 |] |> Data.shuffle |> Data.to_array\n in\n equal ~msg:\"same seed same order\" (array int) d1 d2\n\nlet test_shuffle_different_seed () =\n let a1 =\n Nx.Rng.run ~seed:1 @@ fun () ->\n Data.of_array [| 0; 1; 2; 3; 4; 5; 6; 7 |] |> Data.shuffle |> Data.to_array\n in\n let a2 =\n Nx.Rng.run ~seed:2 @@ fun () ->\n Data.of_array [| 0; 1; 2; 3; 4; 5; 6; 7 |] |> Data.shuffle |> Data.to_array\n in\n is_true ~msg:\"different seed different order\" (a1 <> a2)\n\n(* Consumers *)\n\nlet test_fold () =\n let sum = Data.of_array [| 1; 2; 3; 4 |] |> Data.fold ( + ) 0 in\n equal ~msg:\"fold sum\" int 10 sum\n\nlet test_to_seq () =\n let s = Data.of_array [| 10; 20; 30 |] |> Data.to_seq in\n let a = Array.of_seq s in\n equal ~msg:\"to_seq\" (array int) [| 10; 20; 30 |] a\n\n(* Properties *)\n\nlet test_reset () =\n let d = Data.of_array [| 1; 2; 3 |] in\n let a1 = Data.to_array d in\n Data.reset d;\n let a2 = Data.to_array d in\n equal ~msg:\"reset re-iterates\" (array int) a1 a2\n\nlet test_length () =\n let d = Data.of_array [| 1; 2; 3 |] in\n equal ~msg:\"known length\" (option int) (Some 3) (Data.length d);\n let d2 = Data.map (fun x -> x + 1) d in\n equal ~msg:\"map preserves length\" (option int) (Some 3) (Data.length d2)\n\n(* Utilities *)\n\nlet test_stack_batch () =\n let tensors =\n [|\n Nx.create dtype [| 2 |] [| 1.0; 2.0 |];\n Nx.create dtype [| 2 |] [| 3.0; 4.0 |];\n Nx.create dtype [| 2 |] [| 5.0; 6.0 |];\n |]\n in\n let stacked = Data.stack_batch tensors in\n equal ~msg:\"shape\" (list int) [ 3; 2 ] (Array.to_list (Nx.shape stacked));\n equal ~msg:\"value\" (float 1e-6) 3.0 (Nx.item [ 1; 0 ] stacked)\n\nlet () =\n run \"Kaun.Data\"\n [\n group \"constructors\"\n [\n test \"of_array\" test_of_array;\n test \"of_fn\" test_of_fn;\n test \"of_fn negative\" test_of_fn_negative;\n test \"of_tensor\" test_of_tensor;\n test \"of_tensors\" test_of_tensors;\n test \"of_tensors mismatch\" test_of_tensors_mismatch;\n ];\n group \"transformers\"\n [\n test \"map\" test_map;\n test \"batch\" test_batch;\n test \"batch drop_last\" test_batch_drop_last;\n test \"batch invalid size\" test_batch_invalid_size;\n test \"map_batch\" test_map_batch;\n test \"shuffle deterministic\" test_shuffle_deterministic;\n test \"shuffle different seed\" test_shuffle_different_seed;\n ];\n group \"consumers\" [ test \"fold\" test_fold; test \"to_seq\" test_to_seq ];\n group \"properties\" [ test \"reset\" test_reset; test \"length\" test_length ];\n group \"utilities\" [ test \"stack_batch\" test_stack_batch ];\n ]\n"
},
{
"path": "packages/kaun/test/test_fn.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Fn = Kaun.Fn\n\nlet flatten_f32 t = Nx.to_array (Nx.reshape [| -1 |] (Nx.cast Nx.float32 t))\nlet check_shape msg expected t = equal ~msg (array int) expected (Nx.shape t)\n\nlet check_values msg expected t =\n let actual = flatten_f32 t in\n let n = Array.length expected in\n if Array.length actual <> n then\n failf \"%s: expected %d elements, got %d\" msg n (Array.length actual);\n for i = 0 to n - 1 do\n equal\n ~msg:(Printf.sprintf \"%s[%d]\" msg i)\n (float 1e-4) expected.(i) actual.(i)\n done\n\n(* conv1d *)\n\nlet test_conv1d_basic () =\n let x = Nx.create Nx.float32 [| 1; 1; 5 |] [| 1.; 2.; 3.; 4.; 5. |] in\n let w = Nx.create Nx.float32 [| 1; 1; 3 |] [| 1.; 1.; 1. |] in\n let result = Fn.conv1d x w in\n check_shape \"conv1d basic shape\" [| 1; 1; 3 |] result;\n check_values \"conv1d basic\" [| 6.; 9.; 12. |] result\n\nlet test_conv1d_same_padding () =\n let x = Nx.create Nx.float32 [| 1; 1; 5 |] [| 1.; 2.; 3.; 4.; 5. |] in\n let w = Nx.create Nx.float32 [| 1; 1; 3 |] [| 1.; 1.; 1. |] in\n let result = Fn.conv1d ~padding:`Same x w in\n check_shape \"conv1d same shape\" [| 1; 1; 5 |] result;\n check_values \"conv1d same\" [| 3.; 6.; 9.; 12.; 9. |] result\n\nlet test_conv1d_stride () =\n let x =\n Nx.create Nx.float32 [| 1; 1; 8 |] [| 1.; 2.; 3.; 4.; 5.; 6.; 7.; 8. |]\n in\n let w = Nx.create Nx.float32 [| 1; 1; 3 |] [| 1.; 1.; 1. |] in\n let result = Fn.conv1d ~stride:2 x w in\n check_shape \"conv1d stride shape\" [| 1; 1; 3 |] result;\n check_values \"conv1d stride\" [| 6.; 12.; 18. |] result\n\nlet test_conv1d_dilation () =\n let x = Nx.create Nx.float32 [| 1; 1; 7 |] [| 1.; 2.; 3.; 4.; 5.; 6.; 7. |] in\n let w = Nx.create Nx.float32 [| 1; 1; 3 |] [| 1.; 0.; 1. |] in\n let result = Fn.conv1d ~dilation:2 x w in\n check_shape \"conv1d dilation shape\" [| 1; 1; 3 |] result;\n (* kernel [1;0;1] with dilation=2 picks (i, i+2, i+4): 1+5=6, 2+6=8, 3+7=10 *)\n check_values \"conv1d dilation\" [| 6.; 8.; 10. |] result\n\nlet test_conv1d_bias () =\n let x = Nx.create Nx.float32 [| 1; 1; 3 |] [| 1.; 2.; 3. |] in\n let w = Nx.create Nx.float32 [| 1; 1; 2 |] [| 1.; 1. |] in\n let bias = Nx.create Nx.float32 [| 1 |] [| 10. |] in\n let result = Fn.conv1d ~bias x w in\n check_values \"conv1d bias\" [| 13.; 15. |] result\n\nlet test_conv1d_groups () =\n let x = Nx.create Nx.float32 [| 1; 4; 4 |] (Array.init 16 float_of_int) in\n let w =\n Nx.create Nx.float32 [| 2; 2; 2 |] [| 1.; 1.; 1.; 1.; 1.; 1.; 1.; 1. |]\n in\n let result = Fn.conv1d ~groups:2 x w in\n check_shape \"conv1d groups shape\" [| 1; 2; 3 |] result\n\n(* conv2d *)\n\nlet test_conv2d_basic () =\n let x = Nx.create Nx.float32 [| 1; 1; 4; 4 |] (Array.init 16 float_of_int) in\n let w = Nx.create Nx.float32 [| 1; 1; 3; 3 |] (Array.make 9 1.0) in\n let result = Fn.conv2d x w in\n check_shape \"conv2d basic shape\" [| 1; 1; 2; 2 |] result;\n check_values \"conv2d basic\" [| 45.; 54.; 81.; 90. |] result\n\nlet test_conv2d_same_padding () =\n let x =\n Nx.create Nx.float32 [| 1; 1; 3; 3 |]\n [| 1.; 2.; 3.; 4.; 5.; 6.; 7.; 8.; 9. |]\n in\n let w = Nx.create Nx.float32 [| 1; 1; 2; 2 |] [| 1.; 1.; 1.; 1. |] in\n let result = Fn.conv2d ~padding:`Same x w in\n check_shape \"conv2d same shape\" [| 1; 1; 3; 3 |] result\n\nlet test_conv2d_stride () =\n let x = Nx.create Nx.float32 [| 1; 1; 5; 5 |] (Array.init 25 float_of_int) in\n let w = Nx.create Nx.float32 [| 1; 1; 3; 3 |] (Array.make 9 1.0) in\n let result = Fn.conv2d ~stride:(2, 2) x w in\n check_shape \"conv2d stride shape\" [| 1; 1; 2; 2 |] result;\n check_values \"conv2d stride\" [| 54.; 72.; 144.; 162. |] result\n\nlet test_conv2d_dilation () =\n let x = Nx.create Nx.float32 [| 1; 1; 5; 5 |] (Array.init 25 float_of_int) in\n let w =\n Nx.create Nx.float32 [| 1; 1; 3; 3 |]\n [| 1.; 0.; 0.; 0.; 0.; 0.; 0.; 0.; 1. |]\n in\n let result = Fn.conv2d ~dilation:(2, 2) x w in\n check_shape \"conv2d dilation shape\" [| 1; 1; 1; 1 |] result;\n check_values \"conv2d dilation\" [| 24. |] result\n\nlet test_conv2d_multi_channel () =\n let x = Nx.create Nx.float32 [| 1; 3; 4; 4 |] (Array.init 48 float_of_int) in\n let w = Nx.create Nx.float32 [| 2; 3; 3; 3 |] (Array.make 54 1.0) in\n let result = Fn.conv2d x w in\n check_shape \"conv2d multi-channel shape\" [| 1; 2; 2; 2 |] result\n\nlet test_conv2d_groups () =\n let x = Nx.create Nx.float32 [| 1; 4; 6; 6 |] (Array.init 144 float_of_int) in\n let w = Nx.create Nx.float32 [| 4; 2; 2; 2 |] (Array.make 32 1.0) in\n let result = Fn.conv2d ~groups:2 x w in\n check_shape \"conv2d groups shape\" [| 1; 4; 5; 5 |] result\n\nlet test_conv2d_bias () =\n let x = Nx.create Nx.float32 [| 1; 1; 3; 3 |] (Array.make 9 1.0) in\n let w = Nx.create Nx.float32 [| 2; 1; 2; 2 |] (Array.make 8 1.0) in\n let bias = Nx.create Nx.float32 [| 2 |] [| 10.; 20. |] in\n let result = Fn.conv2d ~bias x w in\n check_shape \"conv2d bias shape\" [| 1; 2; 2; 2 |] result;\n (* Each 2x2 window of ones with all-ones kernel = 4.0, + bias *)\n check_values \"conv2d bias\" [| 14.; 14.; 14.; 14.; 24.; 24.; 24.; 24. |] result\n\n(* max_pool1d *)\n\nlet test_max_pool1d_basic () =\n let x = Nx.create Nx.float32 [| 1; 1; 6 |] [| 1.; 3.; 2.; 5.; 4.; 6. |] in\n let result = Fn.max_pool1d ~kernel_size:2 ~stride:2 x in\n check_shape \"max_pool1d shape\" [| 1; 1; 3 |] result;\n check_values \"max_pool1d\" [| 3.; 5.; 6. |] result\n\nlet test_max_pool1d_same_padding () =\n let x = Nx.create Nx.float32 [| 1; 1; 5 |] [| 1.; 3.; 2.; 5.; 4. |] in\n let result = Fn.max_pool1d ~kernel_size:3 ~stride:1 ~padding:`Same x in\n check_shape \"max_pool1d same shape\" [| 1; 1; 5 |] result\n\n(* max_pool2d *)\n\nlet test_max_pool2d_basic () =\n let x = Nx.create Nx.float32 [| 1; 1; 4; 4 |] (Array.init 16 float_of_int) in\n let result = Fn.max_pool2d ~kernel_size:(2, 2) ~stride:(2, 2) x in\n check_shape \"max_pool2d shape\" [| 1; 1; 2; 2 |] result;\n check_values \"max_pool2d\" [| 5.; 7.; 13.; 15. |] result\n\nlet test_max_pool2d_stride1 () =\n let x =\n Nx.create Nx.float32 [| 1; 1; 3; 3 |]\n [| 1.; 2.; 3.; 4.; 5.; 6.; 7.; 8.; 9. |]\n in\n let result = Fn.max_pool2d ~kernel_size:(2, 2) ~stride:(1, 1) x in\n check_shape \"max_pool2d stride1 shape\" [| 1; 1; 2; 2 |] result;\n check_values \"max_pool2d stride1\" [| 5.; 6.; 8.; 9. |] result\n\n(* avg_pool1d *)\n\nlet test_avg_pool1d_basic () =\n let x = Nx.create Nx.float32 [| 1; 1; 6 |] [| 1.; 2.; 3.; 4.; 5.; 6. |] in\n let result = Fn.avg_pool1d ~kernel_size:2 ~stride:2 x in\n check_shape \"avg_pool1d shape\" [| 1; 1; 3 |] result;\n check_values \"avg_pool1d\" [| 1.5; 3.5; 5.5 |] result\n\n(* avg_pool2d *)\n\nlet test_avg_pool2d_basic () =\n let x = Nx.create Nx.float32 [| 1; 1; 4; 4 |] (Array.init 16 float_of_int) in\n let result = Fn.avg_pool2d ~kernel_size:(2, 2) ~stride:(2, 2) x in\n check_shape \"avg_pool2d shape\" [| 1; 1; 2; 2 |] result;\n check_values \"avg_pool2d\" [| 2.5; 4.5; 10.5; 12.5 |] result\n\nlet test_avg_pool2d_same_padding () =\n let x =\n Nx.create Nx.float32 [| 1; 1; 3; 3 |]\n [| 1.; 2.; 3.; 4.; 5.; 6.; 7.; 8.; 9. |]\n in\n let result =\n Fn.avg_pool2d ~kernel_size:(2, 2) ~stride:(2, 2) ~padding:`Same x\n in\n check_shape \"avg_pool2d same shape\" [| 1; 1; 2; 2 |] result\n\n(* Gradient tests *)\n\nlet eps = 1e-4\n\nlet check_rune ~eps msg expected actual =\n let xs = flatten_f32 expected in\n let ys = flatten_f32 actual in\n let n = Array.length xs in\n if Array.length ys <> n then\n failf \"%s: shape mismatch: expected %d elts, got %d\" msg n (Array.length ys);\n for i = 0 to n - 1 do\n equal ~msg:(Printf.sprintf \"%s[%d]\" msg i) (float eps) xs.(i) ys.(i)\n done\n\nlet test_grad_conv2d () =\n (* conv2d is correlation (no kernel flip), unlike the old Nx convolve2d *)\n let x =\n Nx.create Nx.float32 [| 1; 1; 4; 4 |]\n (Array.init 16 (fun i -> float_of_int (i + 1)))\n in\n let w = Nx.create Nx.float32 [| 1; 1; 2; 2 |] [| 1.; 0.; 0.; 1. |] in\n (* grad w.r.t. input: sum(conv2d(x, w)) → each input pixel's grad is how many\n output windows include it, weighted by the kernel value at that position.\n For a 2x2 kernel [1,0;0,1] on 4x4 input with Valid padding → 3x3 output.\n JAX: jax.grad(lambda x: jnp.sum(jax.lax.conv(x, w, (1,1), 'VALID')))(x) *)\n let f_x x = Nx.sum (Fn.conv2d x w) in\n let grad_x = Rune.grad f_x x in\n let expected_x =\n Nx.create Nx.float32 [| 1; 1; 4; 4 |]\n [| 1.; 1.; 1.; 0.; 1.; 2.; 2.; 1.; 1.; 2.; 2.; 1.; 0.; 1.; 1.; 1. |]\n in\n check_rune ~eps \"conv2d dx\" expected_x grad_x;\n (* grad w.r.t. kernel *)\n let f_w w = Nx.sum (Fn.conv2d x w) in\n let grad_w = Rune.grad f_w w in\n (* For correlation: dL/dw[i,j] = sum of x values at positions covered by\n w[i,j] across all output windows. w[0,0] covers x[0..2,0..2], w[0,1] covers\n x[0..2,1..3], etc. *)\n let expected_w =\n Nx.create Nx.float32 [| 1; 1; 2; 2 |] [| 54.; 63.; 90.; 99. |]\n in\n check_rune ~eps \"conv2d dw\" expected_w grad_w\n\nlet test_grad_avg_pool2d () =\n let x =\n Nx.create Nx.float32 [| 1; 1; 4; 4 |]\n (Array.init 16 (fun i -> float_of_int (i + 1)))\n in\n (* Non-overlapping 2x2 avg pool: each output = mean of 4 inputs. grad of\n sum(avg_pool) = 0.25 everywhere (each input contributes to exactly one\n output, scaled by 1/4) *)\n let f x = Nx.sum (Fn.avg_pool2d ~kernel_size:(2, 2) ~stride:(2, 2) x) in\n let grad_x = Rune.grad f x in\n let expected = Nx.full Nx.float32 [| 1; 1; 4; 4 |] 0.25 in\n check_rune ~eps \"avg_pool2d dx\" expected grad_x\n\nlet test_grad_avg_pool2d_overlapping () =\n let x =\n Nx.create Nx.float32 [| 1; 1; 4; 4 |]\n (Array.init 16 (fun i -> float_of_int (i + 1)))\n in\n (* Overlapping 2x2 avg pool with stride 1: 3x3 output. Each output window\n contributes 0.25 per input pixel it covers. Corner pixels appear in 1\n window, edge in 2, interior in 4. *)\n let f x = Nx.sum (Fn.avg_pool2d ~kernel_size:(2, 2) ~stride:(1, 1) x) in\n let grad_x = Rune.grad f x in\n let expected =\n Nx.create Nx.float32 [| 1; 1; 4; 4 |]\n [|\n 0.25;\n 0.5;\n 0.5;\n 0.25;\n 0.5;\n 1.0;\n 1.0;\n 0.5;\n 0.5;\n 1.0;\n 1.0;\n 0.5;\n 0.25;\n 0.5;\n 0.5;\n 0.25;\n |]\n in\n check_rune ~eps \"avg_pool2d overlapping dx\" expected grad_x\n\nlet () =\n run \"Kaun.Fn\"\n [\n group \"conv1d\"\n [\n test \"basic\" test_conv1d_basic;\n test \"same padding\" test_conv1d_same_padding;\n test \"stride\" test_conv1d_stride;\n test \"dilation\" test_conv1d_dilation;\n test \"bias\" test_conv1d_bias;\n test \"groups\" test_conv1d_groups;\n ];\n group \"conv2d\"\n [\n test \"basic\" test_conv2d_basic;\n test \"same padding\" test_conv2d_same_padding;\n test \"stride\" test_conv2d_stride;\n test \"dilation\" test_conv2d_dilation;\n test \"multi-channel\" test_conv2d_multi_channel;\n test \"groups\" test_conv2d_groups;\n test \"bias\" test_conv2d_bias;\n ];\n group \"max_pool\"\n [\n test \"1d basic\" test_max_pool1d_basic;\n test \"1d same padding\" test_max_pool1d_same_padding;\n test \"2d basic\" test_max_pool2d_basic;\n test \"2d stride 1\" test_max_pool2d_stride1;\n ];\n group \"avg_pool\"\n [\n test \"1d basic\" test_avg_pool1d_basic;\n test \"2d basic\" test_avg_pool2d_basic;\n test \"2d same padding\" test_avg_pool2d_same_padding;\n ];\n group \"gradients\"\n [\n test \"conv2d\" test_grad_conv2d;\n test \"avg_pool2d\" test_grad_avg_pool2d;\n test \"avg_pool2d overlapping\" test_grad_avg_pool2d_overlapping;\n ];\n ]\n"
},
{
"path": "packages/kaun/test/test_grad.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Grad = Kaun.Grad\nmodule Ptree = Kaun.Ptree\n\nlet string_contains s sub =\n let slen = String.length s in\n let sub_len = String.length sub in\n let rec loop i =\n if i + sub_len > slen then false\n else if String.sub s i sub_len = sub then true\n else loop (i + 1)\n in\n if sub_len = 0 then true else loop 0\n\nlet raises_invalid_arg_contains needle f =\n raises_match\n (fun exn ->\n match exn with\n | Invalid_argument msg -> string_contains msg needle\n | _ -> false)\n f\n\n(* f(x) = 0.5 * sum(x^2), gradient = x *)\nlet test_scalar_quadratic () =\n let x = Nx.create Nx.float32 [| 2 |] [| 3.0; -4.0 |] in\n let params = Ptree.tensor x in\n let loss, grads =\n Grad.value_and_grad\n (fun p ->\n let (Ptree.P t) = Ptree.as_tensor_exn p in\n let t = Ptree.Tensor.to_typed_exn Nx.float32 (Ptree.P t) in\n Nx.mul (Nx.scalar Nx.float32 0.5) (Nx.sum (Nx.mul t t)))\n params\n in\n equal ~msg:\"loss value\" (float 1e-6) 12.5 (Nx.item [] loss);\n let g = Ptree.Tensor.to_typed_exn Nx.float32 (Ptree.as_tensor_exn grads) in\n equal ~msg:\"grad[0]\" (float 1e-5) 3.0 (Nx.item [ 0 ] g);\n equal ~msg:\"grad[1]\" (float 1e-5) (-4.0) (Nx.item [ 1 ] g)\n\n(* f(w, b) = sum(w * x + b), dw = x, db = ones *)\nlet test_multi_leaf_dict () =\n let w = Nx.create Nx.float32 [| 3 |] [| 1.0; 2.0; 3.0 |] in\n let b = Nx.create Nx.float32 [| 3 |] [| 0.1; 0.2; 0.3 |] in\n let x = Nx.create Nx.float32 [| 3 |] [| 4.0; 5.0; 6.0 |] in\n let params = Ptree.dict [ (\"w\", Ptree.tensor w); (\"b\", Ptree.tensor b) ] in\n let loss, grads =\n Grad.value_and_grad\n (fun p ->\n let fields = Ptree.Dict.fields_exn p in\n let w = Ptree.Dict.get_tensor_exn fields ~name:\"w\" Nx.float32 in\n let b = Ptree.Dict.get_tensor_exn fields ~name:\"b\" Nx.float32 in\n Nx.sum (Nx.add (Nx.mul w x) b))\n params\n in\n equal ~msg:\"loss value\" (float 1e-4) 32.6 (Nx.item [] loss);\n let grad_fields = Ptree.Dict.fields_exn grads in\n let gw = Ptree.Dict.get_tensor_exn grad_fields ~name:\"w\" Nx.float32 in\n let gb = Ptree.Dict.get_tensor_exn grad_fields ~name:\"b\" Nx.float32 in\n equal ~msg:\"dw[0]\" (float 1e-5) 4.0 (Nx.item [ 0 ] gw);\n equal ~msg:\"dw[1]\" (float 1e-5) 5.0 (Nx.item [ 1 ] gw);\n equal ~msg:\"dw[2]\" (float 1e-5) 6.0 (Nx.item [ 2 ] gw);\n equal ~msg:\"db[0]\" (float 1e-5) 1.0 (Nx.item [ 0 ] gb);\n equal ~msg:\"db[1]\" (float 1e-5) 1.0 (Nx.item [ 1 ] gb);\n equal ~msg:\"db[2]\" (float 1e-5) 1.0 (Nx.item [ 2 ] gb)\n\nlet test_nested_tree () =\n let a = Nx.create Nx.float32 [| 2 |] [| 1.0; 2.0 |] in\n let b = Nx.create Nx.float32 [| 2 |] [| 3.0; 4.0 |] in\n let params =\n Ptree.dict\n [\n (\"layer1\", Ptree.dict [ (\"w\", Ptree.tensor a) ]);\n (\"layer2\", Ptree.dict [ (\"w\", Ptree.tensor b) ]);\n ]\n in\n let _loss, grads =\n Grad.value_and_grad\n (fun p ->\n let f1 = Ptree.Dict.fields_exn p in\n let l1 = Ptree.Dict.fields_exn (Ptree.Dict.find_exn \"layer1\" f1) in\n let l2 = Ptree.Dict.fields_exn (Ptree.Dict.find_exn \"layer2\" f1) in\n let w1 = Ptree.Dict.get_tensor_exn l1 ~name:\"w\" Nx.float32 in\n let w2 = Ptree.Dict.get_tensor_exn l2 ~name:\"w\" Nx.float32 in\n Nx.add (Nx.sum (Nx.mul w1 w1)) (Nx.sum (Nx.mul w2 w2)))\n params\n in\n let gf = Ptree.Dict.fields_exn grads in\n let gl1 = Ptree.Dict.fields_exn (Ptree.Dict.find_exn \"layer1\" gf) in\n let gl2 = Ptree.Dict.fields_exn (Ptree.Dict.find_exn \"layer2\" gf) in\n let ga = Ptree.Dict.get_tensor_exn gl1 ~name:\"w\" Nx.float32 in\n let gb = Ptree.Dict.get_tensor_exn gl2 ~name:\"w\" Nx.float32 in\n equal ~msg:\"ga[0]\" (float 1e-5) 2.0 (Nx.item [ 0 ] ga);\n equal ~msg:\"ga[1]\" (float 1e-5) 4.0 (Nx.item [ 1 ] ga);\n equal ~msg:\"gb[0]\" (float 1e-5) 6.0 (Nx.item [ 0 ] gb);\n equal ~msg:\"gb[1]\" (float 1e-5) 8.0 (Nx.item [ 1 ] gb)\n\nlet test_value_and_grad_aux () =\n let x = Nx.create Nx.float32 [| 2 |] [| 3.0; 4.0 |] in\n let params = Ptree.tensor x in\n let loss, grads, aux =\n Grad.value_and_grad_aux\n (fun p ->\n let (Ptree.P t) = Ptree.as_tensor_exn p in\n let t = Ptree.Tensor.to_typed_exn Nx.float32 (Ptree.P t) in\n (Nx.sum (Nx.mul t t), Nx.item [ 0 ] t))\n params\n in\n equal ~msg:\"loss value\" (float 1e-6) 25.0 (Nx.item [] loss);\n equal ~msg:\"aux value\" (float 1e-6) 3.0 aux;\n let g = Ptree.Tensor.to_typed_exn Nx.float32 (Ptree.as_tensor_exn grads) in\n equal ~msg:\"grad[0]\" (float 1e-5) 6.0 (Nx.item [ 0 ] g);\n equal ~msg:\"grad[1]\" (float 1e-5) 8.0 (Nx.item [ 1 ] g)\n\nlet test_grad_convenience () =\n let x = Nx.create Nx.float32 [| 2 |] [| 5.0; -3.0 |] in\n let params = Ptree.tensor x in\n let f p =\n let (Ptree.P t) = Ptree.as_tensor_exn p in\n let t = Ptree.Tensor.to_typed_exn Nx.float32 (Ptree.P t) in\n Nx.sum t\n in\n let grads = Grad.grad f params in\n let g = Ptree.Tensor.to_typed_exn Nx.float32 (Ptree.as_tensor_exn grads) in\n equal ~msg:\"grad[0]\" (float 1e-5) 1.0 (Nx.item [ 0 ] g);\n equal ~msg:\"grad[1]\" (float 1e-5) 1.0 (Nx.item [ 1 ] g)\n\nlet test_empty_tree () =\n let params = Ptree.list [] in\n let loss, grads =\n Grad.value_and_grad (fun _p -> Nx.scalar Nx.float32 42.0) params\n in\n equal ~msg:\"empty tree loss\" (float 1e-6) 42.0 (Nx.item [] loss);\n match grads with\n | Ptree.List [] -> ()\n | _ -> fail \"expected empty list gradient\"\n\nlet test_non_float_leaf_error () =\n let params = Ptree.tensor (Nx.zeros Nx.int32 [| 3 |]) in\n raises_invalid_arg_contains \" expected float dtype\" (fun () ->\n ignore (Grad.value_and_grad (fun _p -> Nx.scalar Nx.float32 0.0) params))\n\nlet test_mixed_dtype_error () =\n let params =\n Ptree.dict\n [\n (\"a\", Ptree.tensor (Nx.ones Nx.float16 [| 2 |]));\n (\"b\", Ptree.tensor (Nx.ones Nx.float32 [| 2 |]));\n ]\n in\n raises_invalid_arg_contains \"has dtype/layout\" (fun () ->\n ignore\n (Grad.value_and_grad\n (fun p ->\n let fields = Ptree.Dict.fields_exn p in\n let a = Ptree.Dict.get_tensor_exn fields ~name:\"a\" Nx.float16 in\n let b = Ptree.Dict.get_tensor_exn fields ~name:\"b\" Nx.float32 in\n Nx.add (Nx.sum (Nx.cast Nx.float32 a)) (Nx.sum b))\n params))\n\nlet test_value_and_grad_mixed () =\n let params =\n Ptree.dict\n [\n (\"a\", Ptree.tensor (Nx.ones Nx.float16 [| 2 |]));\n (\"b\", Ptree.tensor (Nx.ones Nx.float32 [| 2 |]));\n ]\n in\n let loss, grads =\n Grad.value_and_grad_mixed\n (fun p ->\n let fields = Ptree.Dict.fields_exn p in\n let a = Ptree.Dict.get_tensor_exn fields ~name:\"a\" Nx.float16 in\n let b = Ptree.Dict.get_tensor_exn fields ~name:\"b\" Nx.float32 in\n Nx.add (Nx.sum (Nx.cast Nx.float32 a)) (Nx.sum b))\n params\n in\n equal ~msg:\"loss value\" (float 1e-5) 4.0 (Nx.item [] loss);\n let grad_fields = Ptree.Dict.fields_exn grads in\n let ga = Ptree.Dict.get_tensor_exn grad_fields ~name:\"a\" Nx.float16 in\n let gb = Ptree.Dict.get_tensor_exn grad_fields ~name:\"b\" Nx.float32 in\n equal ~msg:\"ga[0]\" (float 1e-5) 1.0 (Nx.item [ 0 ] (Nx.cast Nx.float32 ga));\n equal ~msg:\"gb[0]\" (float 1e-5) 1.0 (Nx.item [ 0 ] gb)\n\nlet test_structure_preserved () =\n let params =\n Ptree.list\n [\n Ptree.tensor (Nx.ones Nx.float32 [| 2 |]);\n Ptree.tensor (Nx.ones Nx.float32 [| 3 |]);\n ]\n in\n let grads =\n Grad.grad\n (fun p ->\n let items = Ptree.List.items_exn p in\n let t0 =\n Ptree.Tensor.to_typed_exn Nx.float32\n (Ptree.as_tensor_exn (List.nth items 0))\n in\n let t1 =\n Ptree.Tensor.to_typed_exn Nx.float32\n (Ptree.as_tensor_exn (List.nth items 1))\n in\n Nx.add (Nx.sum t0) (Nx.sum t1))\n params\n in\n let items = Ptree.List.items_exn grads in\n equal ~msg:\"list length\" int 2 (List.length items);\n let g0 =\n Ptree.Tensor.to_typed_exn Nx.float32\n (Ptree.as_tensor_exn (List.nth items 0))\n in\n let g1 =\n Ptree.Tensor.to_typed_exn Nx.float32\n (Ptree.as_tensor_exn (List.nth items 1))\n in\n equal ~msg:\"g0 shape\" (list int) [ 2 ] (Array.to_list (Nx.shape g0));\n equal ~msg:\"g1 shape\" (list int) [ 3 ] (Array.to_list (Nx.shape g1))\n\nlet () =\n run \"Kaun.Grad\"\n [\n group \"value_and_grad\"\n [\n test \"scalar quadratic\" test_scalar_quadratic;\n test \"multi-leaf dict\" test_multi_leaf_dict;\n test \"nested tree\" test_nested_tree;\n test \"structure preserved\" test_structure_preserved;\n test \"value_and_grad_aux\" test_value_and_grad_aux;\n test \"value_and_grad_mixed\" test_value_and_grad_mixed;\n ];\n group \"grad\" [ test \"grad convenience\" test_grad_convenience ];\n group \"edge cases\"\n [\n test \"empty tree\" test_empty_tree;\n test \"non-float leaf error\" test_non_float_leaf_error;\n test \"mixed dtype error\" test_mixed_dtype_error;\n ];\n ]\n"
},
{
"path": "packages/kaun/test/test_init.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Init = Kaun.Init\n\nlet string_contains s sub =\n let slen = String.length s in\n let sub_len = String.length sub in\n let rec loop i =\n if i + sub_len > slen then false\n else if String.sub s i sub_len = sub then true\n else loop (i + 1)\n in\n if sub_len = 0 then true else loop 0\n\nlet raises_invalid_arg_contains needle f =\n raises_match\n (fun exn ->\n match exn with\n | Invalid_argument msg -> string_contains msg needle\n | _ -> false)\n f\n\nlet flatten_f32 t = Nx.to_array (Nx.reshape [| -1 |] (Nx.cast Nx.float32 t))\n\nlet tensor_all pred t =\n let a = flatten_f32 t in\n Array.for_all pred a\n\nlet tensor_stats t =\n let a = flatten_f32 t in\n let n = Array.length a in\n let sum = ref 0.0 in\n for i = 0 to n - 1 do\n sum := !sum +. a.(i)\n done;\n let mean = !sum /. float_of_int n in\n let sq = ref 0.0 in\n for i = 0 to n - 1 do\n let d = a.(i) -. mean in\n sq := !sq +. (d *. d)\n done;\n let variance = !sq /. float_of_int n in\n (mean, variance)\n\nlet compute_fans shape ~in_axis ~out_axis =\n let rank = Array.length shape in\n if rank = 0 then (1, 1)\n else if rank = 1 then (shape.(0), shape.(0))\n else\n let normalize_axis axis = if axis < 0 then rank + axis else axis in\n let in_axis = normalize_axis in_axis in\n let out_axis = normalize_axis out_axis in\n let fan_in = shape.(in_axis) in\n let fan_out = shape.(out_axis) in\n let receptive = ref 1 in\n for i = 0 to rank - 1 do\n if i <> in_axis && i <> out_axis then receptive := !receptive * shape.(i)\n done;\n (fan_in * !receptive, fan_out * !receptive)\n\nlet expected_variance ~scale ~mode ~fan_in ~fan_out =\n let n =\n match mode with\n | `Fan_in -> float_of_int fan_in\n | `Fan_out -> float_of_int fan_out\n | `Fan_avg -> float_of_int (fan_in + fan_out) /. 2.0\n in\n scale /. n\n\nlet uniform_limit variance = sqrt (3.0 *. variance)\n\nlet test_constants () =\n Nx.Rng.run ~seed:0 @@ fun () ->\n let shape = [| 11; 13 |] in\n let zeros = Init.zeros.f shape Nx.float32 in\n equal ~msg:\"zeros\" bool true (tensor_all (fun x -> x = 0.0) zeros);\n let ones = Init.ones.f shape Nx.float32 in\n equal ~msg:\"ones\" bool true (tensor_all (fun x -> x = 1.0) ones);\n let c = (Init.constant 3.5).f shape Nx.float32 in\n equal ~msg:\"constant\" bool true (tensor_all (fun x -> x = 3.5) c)\n\nlet test_uniform_range_and_mean () =\n Nx.Rng.run ~seed:1 @@ fun () ->\n let scale = 0.25 in\n let t = (Init.uniform ~scale ()).f [| 120_000 |] Nx.float32 in\n equal ~msg:\"uniform range\" bool true\n (tensor_all (fun x -> x >= 0.0 && x < scale) t);\n let mean, _ = tensor_stats t in\n equal ~msg:\"uniform mean\" (float 8e-3) (scale /. 2.0) mean\n\nlet test_normal_mean_and_variance () =\n Nx.Rng.run ~seed:2 @@ fun () ->\n let stddev = 0.2 in\n let t = (Init.normal ~stddev ()).f [| 140_000 |] Nx.float32 in\n let mean, variance = tensor_stats t in\n equal ~msg:\"normal mean\" (float 6e-3) 0.0 mean;\n equal ~msg:\"normal variance\" (float 8e-3) (stddev *. stddev) variance\n\nlet test_glorot_uniform_bounds () =\n Nx.Rng.run ~seed:3 @@ fun () ->\n let shape = [| 64; 32 |] in\n let fan_in, fan_out = compute_fans shape ~in_axis:(-2) ~out_axis:(-1) in\n let variance = expected_variance ~scale:1.0 ~mode:`Fan_avg ~fan_in ~fan_out in\n let limit = uniform_limit variance in\n let t = (Init.glorot_uniform ()).f shape Nx.float32 in\n equal ~msg:\"glorot_uniform bounds\" bool true\n (tensor_all (fun x -> x >= -.limit && x <= limit) t)\n\nlet test_glorot_normal_variance () =\n Nx.Rng.run ~seed:4 @@ fun () ->\n let shape = [| 960; 480 |] in\n let fan_in, fan_out = compute_fans shape ~in_axis:(-2) ~out_axis:(-1) in\n let expected = expected_variance ~scale:1.0 ~mode:`Fan_avg ~fan_in ~fan_out in\n let t = (Init.glorot_normal ()).f shape Nx.float32 in\n let _, variance = tensor_stats t in\n equal ~msg:\"glorot_normal variance\" (float 3e-4) expected variance\n\nlet test_he_uniform_bounds () =\n Nx.Rng.run ~seed:5 @@ fun () ->\n let shape = [| 128; 64 |] in\n let fan_in, fan_out = compute_fans shape ~in_axis:(-2) ~out_axis:(-1) in\n let variance = expected_variance ~scale:2.0 ~mode:`Fan_in ~fan_in ~fan_out in\n let limit = uniform_limit variance in\n let t = (Init.he_uniform ()).f shape Nx.float32 in\n equal ~msg:\"he_uniform bounds\" bool true\n (tensor_all (fun x -> x >= -.limit && x <= limit) t)\n\nlet test_he_normal_variance () =\n Nx.Rng.run ~seed:6 @@ fun () ->\n let shape = [| 256; 64 |] in\n let fan_in, fan_out = compute_fans shape ~in_axis:(-2) ~out_axis:(-1) in\n let expected = expected_variance ~scale:2.0 ~mode:`Fan_in ~fan_in ~fan_out in\n let t = (Init.he_normal ()).f shape Nx.float32 in\n let _, variance = tensor_stats t in\n equal ~msg:\"he_normal variance\" (float 2e-3) expected variance\n\nlet test_lecun_uniform_bounds () =\n Nx.Rng.run ~seed:7 @@ fun () ->\n let shape = [| 128; 32 |] in\n let fan_in, fan_out = compute_fans shape ~in_axis:(-2) ~out_axis:(-1) in\n let variance = expected_variance ~scale:1.0 ~mode:`Fan_in ~fan_in ~fan_out in\n let limit = uniform_limit variance in\n let t = (Init.lecun_uniform ()).f shape Nx.float32 in\n equal ~msg:\"lecun_uniform bounds\" bool true\n (tensor_all (fun x -> x >= -.limit && x <= limit) t)\n\nlet test_lecun_normal_variance () =\n Nx.Rng.run ~seed:8 @@ fun () ->\n let shape = [| 128; 16 |] in\n let fan_in, fan_out = compute_fans shape ~in_axis:(-2) ~out_axis:(-1) in\n let expected = expected_variance ~scale:1.0 ~mode:`Fan_in ~fan_in ~fan_out in\n let t = (Init.lecun_normal ()).f shape Nx.float32 in\n let _, variance = tensor_stats t in\n equal ~msg:\"lecun_normal variance\" (float 1.5e-3) expected variance\n\nlet test_variance_scaling_axis_override () =\n Nx.Rng.run ~seed:9 @@ fun () ->\n let shape = [| 2; 9; 4 |] in\n let in_axis = 2 in\n let out_axis = 0 in\n let fan_in, fan_out = compute_fans shape ~in_axis ~out_axis in\n let variance = expected_variance ~scale:1.7 ~mode:`Fan_out ~fan_in ~fan_out in\n let limit = uniform_limit variance in\n let init =\n Init.variance_scaling ~scale:1.7 ~mode:`Fan_out ~distribution:`Uniform\n ~in_axis ~out_axis ()\n in\n let t = init.f shape Nx.float32 in\n equal ~msg:\"variance_scaling axis override\" bool true\n (tensor_all (fun x -> x >= -.limit && x <= limit) t)\n\nlet test_validation_errors () =\n raises_invalid_arg_contains \"scale\" (fun () ->\n ignore (Init.uniform ~scale:(-1.0) ()));\n raises_invalid_arg_contains \"stddev\" (fun () ->\n ignore (Init.normal ~stddev:(-0.1) ()));\n raises_invalid_arg_contains \"scale\" (fun () ->\n ignore\n (Init.variance_scaling ~scale:(-1.0) ~mode:`Fan_in\n ~distribution:`Uniform ()));\n let init =\n Init.variance_scaling ~scale:1.0 ~mode:`Fan_avg ~distribution:`Uniform\n ~in_axis:9 ()\n in\n Nx.Rng.run ~seed:10 @@ fun () ->\n raises_invalid_arg_contains \"invalid in axis\" (fun () ->\n ignore (init.f [| 3; 4 |] Nx.float32));\n let zero_fan =\n Init.variance_scaling ~scale:1.0 ~mode:`Fan_in ~distribution:`Uniform ()\n in\n raises_invalid_arg_contains \"non-positive fan\" (fun () ->\n ignore (zero_fan.f [| 0; 4 |] Nx.float32))\n\nlet test_deterministic_same_seed () =\n let init = Init.he_uniform () in\n let shape = [| 64; 64 |] in\n let t0 =\n Nx.Rng.run ~seed:12 @@ fun () -> init.f shape Nx.float32 |> flatten_f32\n in\n let t1 =\n Nx.Rng.run ~seed:12 @@ fun () -> init.f shape Nx.float32 |> flatten_f32\n in\n equal ~msg:\"same seed deterministic\" bool true (t0 = t1)\n\nlet () =\n run \"Kaun.Init\"\n [\n group \"constant\" [ test \"zeros ones constant\" test_constants ];\n group \"random\"\n [\n test \"uniform range and mean\" test_uniform_range_and_mean;\n test \"normal mean and variance\" test_normal_mean_and_variance;\n test \"deterministic same seed\" test_deterministic_same_seed;\n ];\n group \"variance scaling families\"\n [\n test \"glorot uniform bounds\" test_glorot_uniform_bounds;\n test \"glorot normal variance\" test_glorot_normal_variance;\n test \"he uniform bounds\" test_he_uniform_bounds;\n test \"he normal variance\" test_he_normal_variance;\n test \"lecun uniform bounds\" test_lecun_uniform_bounds;\n test \"lecun normal variance\" test_lecun_normal_variance;\n test \"variance scaling axis override\"\n test_variance_scaling_axis_override;\n ];\n group \"validation\" [ test \"invalid arguments\" test_validation_errors ];\n ]\n"
},
{
"path": "packages/kaun/test/test_layer.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Layer = Kaun.Layer\nmodule Ptree = Kaun.Ptree\n\nlet flatten_f32 t = Nx.to_array (Nx.reshape [| -1 |] (Nx.cast Nx.float32 t))\n\nlet tensor_close ~eps ~expected ~actual =\n let xs = flatten_f32 expected in\n let ys = flatten_f32 actual in\n let nx = Array.length xs in\n let ny = Array.length ys in\n if nx <> ny then false\n else\n let ok = ref true in\n for i = 0 to nx - 1 do\n if abs_float (xs.(i) -. ys.(i)) > eps then ok := false\n done;\n !ok\n\nlet apply_out (type a in_elt) (m : (a, float) Layer.t) vars ~training\n (x : (a, in_elt) Nx.t) =\n let y, _ = Layer.apply m vars ~training x in\n y\n\n(* Linear *)\n\nlet test_linear_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.linear ~in_features:4 ~out_features:3 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let w = Ptree.Dict.get_tensor_exn fields ~name:\"weight\" Nx.float32 in\n let b = Ptree.Dict.get_tensor_exn fields ~name:\"bias\" Nx.float32 in\n equal ~msg:\"weight shape\" (list int) [ 4; 3 ] (Array.to_list (Nx.shape w));\n equal ~msg:\"bias shape\" (list int) [ 3 ] (Array.to_list (Nx.shape b))\n\nlet test_linear_forward () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.linear ~in_features:2 ~out_features:3 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 1; 2 |] in\n let y = apply_out m vars ~training:false x in\n equal ~msg:\"output shape\" (list int) [ 1; 3 ] (Array.to_list (Nx.shape y))\n\nlet test_linear_manual_params () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.linear ~in_features:2 ~out_features:2 () in\n let w = Nx.create Nx.float32 [| 2; 2 |] [| 1.0; 0.0; 0.0; 1.0 |] in\n let b = Nx.create Nx.float32 [| 2 |] [| 0.5; -0.5 |] in\n let params =\n Ptree.dict [ (\"weight\", Ptree.tensor w); (\"bias\", Ptree.tensor b) ]\n in\n let vars =\n Layer.init m ~dtype:Nx.float32 |> fun vars -> Layer.with_params vars params\n in\n let x = Nx.create Nx.float32 [| 1; 2 |] [| 3.0; 4.0 |] in\n let y = apply_out m vars ~training:false x in\n let expected = Nx.create Nx.float32 [| 1; 2 |] [| 3.5; 3.5 |] in\n equal ~msg:\"linear identity + bias\" bool true\n (tensor_close ~eps:1e-6 ~expected ~actual:y)\n\n(* Normalization *)\n\nlet test_layer_norm_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.layer_norm ~dim:8 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let gamma = Ptree.Dict.get_tensor_exn fields ~name:\"gamma\" Nx.float32 in\n let beta = Ptree.Dict.get_tensor_exn fields ~name:\"beta\" Nx.float32 in\n equal ~msg:\"gamma shape\" (list int) [ 8 ] (Array.to_list (Nx.shape gamma));\n equal ~msg:\"beta shape\" (list int) [ 8 ] (Array.to_list (Nx.shape beta));\n equal ~msg:\"gamma values\" bool true\n (Array.for_all (fun x -> x = 1.0) (flatten_f32 gamma));\n equal ~msg:\"beta values\" bool true\n (Array.for_all (fun x -> x = 0.0) (flatten_f32 beta))\n\nlet test_layer_norm_forward () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.layer_norm ~dim:4 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x =\n Nx.create Nx.float32 [| 2; 4 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0; 7.0; 8.0 |]\n in\n let y = apply_out m vars ~training:false x in\n equal ~msg:\"output shape\" (list int) [ 2; 4 ] (Array.to_list (Nx.shape y))\n\nlet test_rms_norm_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.rms_norm ~dim:6 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let scale = Ptree.Dict.get_tensor_exn fields ~name:\"scale\" Nx.float32 in\n equal ~msg:\"scale shape\" (list int) [ 6 ] (Array.to_list (Nx.shape scale));\n equal ~msg:\"scale values\" bool true\n (Array.for_all (fun x -> x = 1.0) (flatten_f32 scale))\n\nlet test_batch_norm_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.batch_norm ~num_features:3 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let param_fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let state_fields = Ptree.Dict.fields_exn (Layer.state vars) in\n let scale = Ptree.Dict.get_tensor_exn param_fields ~name:\"scale\" Nx.float32 in\n let bias = Ptree.Dict.get_tensor_exn param_fields ~name:\"bias\" Nx.float32 in\n let running_mean =\n Ptree.Dict.get_tensor_exn state_fields ~name:\"running_mean\" Nx.float32\n in\n let running_var =\n Ptree.Dict.get_tensor_exn state_fields ~name:\"running_var\" Nx.float32\n in\n equal ~msg:\"scale shape\" (list int) [ 3 ] (Array.to_list (Nx.shape scale));\n equal ~msg:\"bias shape\" (list int) [ 3 ] (Array.to_list (Nx.shape bias));\n equal ~msg:\"running_mean shape\" (list int) [ 3 ]\n (Array.to_list (Nx.shape running_mean));\n equal ~msg:\"running_var shape\" (list int) [ 3 ]\n (Array.to_list (Nx.shape running_var))\n\nlet test_batch_norm_rank3_axes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.batch_norm ~num_features:3 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let param_fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let scale = Ptree.Dict.get_tensor_exn param_fields ~name:\"scale\" Nx.float32 in\n let bias = Ptree.Dict.get_tensor_exn param_fields ~name:\"bias\" Nx.float32 in\n let x =\n Nx.create Nx.float32 [| 2; 3; 4 |]\n [|\n 1.0;\n 2.0;\n 3.0;\n 4.0;\n 5.0;\n 6.0;\n 7.0;\n 8.0;\n 9.0;\n 10.0;\n 11.0;\n 12.0;\n 2.0;\n 4.0;\n 6.0;\n 8.0;\n 1.0;\n 3.0;\n 5.0;\n 7.0;\n 0.5;\n 1.5;\n 2.5;\n 3.5;\n |]\n in\n let y, _ = Layer.apply m vars ~training:true x in\n let expected = Kaun.Fn.batch_norm ~axes:[ 0; 2 ] ~scale ~bias x in\n equal ~msg:\"batch_norm rank3 uses [0;2] axes\" bool true\n (tensor_close ~eps:1e-6 ~expected ~actual:y)\n\nlet test_batch_norm_running_stats_eval () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.batch_norm ~num_features:2 () in\n let vars0 = Layer.init m ~dtype:Nx.float32 in\n let x_train = Nx.create Nx.float32 [| 2; 2 |] [| 1.0; 2.0; 3.0; 4.0 |] in\n let _y_train, vars1 = Layer.apply m vars0 ~training:true x_train in\n let param_fields = Ptree.Dict.fields_exn (Layer.params vars1) in\n let state_fields = Ptree.Dict.fields_exn (Layer.state vars1) in\n let scale = Ptree.Dict.get_tensor_exn param_fields ~name:\"scale\" Nx.float32 in\n let bias = Ptree.Dict.get_tensor_exn param_fields ~name:\"bias\" Nx.float32 in\n let running_mean =\n Ptree.Dict.get_tensor_exn state_fields ~name:\"running_mean\" Nx.float32\n in\n let running_var =\n Ptree.Dict.get_tensor_exn state_fields ~name:\"running_var\" Nx.float32\n in\n let x_eval = Nx.create Nx.float32 [| 2; 2 |] [| 10.0; 20.0; 30.0; 40.0 |] in\n let y_eval, vars2 = Layer.apply m vars1 ~training:false x_eval in\n let expected =\n Nx.standardize ~axes:[ 0 ] ~mean:running_mean ~variance:running_var x_eval\n |> fun z ->\n Nx.add (Nx.mul z (Nx.reshape [| 1; 2 |] scale)) (Nx.reshape [| 1; 2 |] bias)\n in\n equal ~msg:\"batch_norm eval uses running stats\" bool true\n (tensor_close ~eps:1e-6 ~expected ~actual:y_eval);\n let state_fields2 = Ptree.Dict.fields_exn (Layer.state vars2) in\n let running_mean2 =\n Ptree.Dict.get_tensor_exn state_fields2 ~name:\"running_mean\" Nx.float32\n in\n let running_var2 =\n Ptree.Dict.get_tensor_exn state_fields2 ~name:\"running_var\" Nx.float32\n in\n equal ~msg:\"batch_norm eval keeps running_mean\" bool true\n (tensor_close ~eps:1e-6 ~expected:running_mean ~actual:running_mean2);\n equal ~msg:\"batch_norm eval keeps running_var\" bool true\n (tensor_close ~eps:1e-6 ~expected:running_var ~actual:running_var2)\n\nlet test_batch_norm_eval_affine_rank3 () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.batch_norm ~num_features:3 () in\n let scale = Nx.create Nx.float32 [| 3 |] [| 2.0; 3.0; 4.0 |] in\n let bias = Nx.create Nx.float32 [| 3 |] [| 10.0; 20.0; 30.0 |] in\n let running_mean = Nx.create Nx.float32 [| 3 |] [| 1.0; 2.0; 3.0 |] in\n let running_var = Nx.create Nx.float32 [| 3 |] [| 4.0; 9.0; 16.0 |] in\n let vars =\n Layer.init m ~dtype:Nx.float32 |> fun vars ->\n Layer.with_params vars\n (Ptree.dict\n [ (\"scale\", Ptree.tensor scale); (\"bias\", Ptree.tensor bias) ])\n |> fun vars ->\n Layer.with_state vars\n (Ptree.dict\n [\n (\"running_mean\", Ptree.tensor running_mean);\n (\"running_var\", Ptree.tensor running_var);\n ])\n in\n let x =\n Nx.create Nx.float32 [| 1; 3; 2 |] [| 1.0; 5.0; 2.0; 8.0; 3.0; 11.0 |]\n in\n let y, _ = Layer.apply m vars ~training:false x in\n let expected =\n Nx.standardize ~axes:[ 0; 2 ] ~mean:running_mean ~variance:running_var x\n |> fun z ->\n Nx.add\n (Nx.mul z (Nx.reshape [| 1; 3; 1 |] scale))\n (Nx.reshape [| 1; 3; 1 |] bias)\n in\n equal ~msg:\"batch_norm eval rank3 applies affine\" bool true\n (tensor_close ~eps:1e-6 ~expected ~actual:y)\n\n(* Embedding *)\n\nlet test_embedding_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.embedding ~vocab_size:100 ~embed_dim:16 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let emb = Ptree.Dict.get_tensor_exn fields ~name:\"embedding\" Nx.float32 in\n equal ~msg:\"embedding shape\" (list int) [ 100; 16 ]\n (Array.to_list (Nx.shape emb))\n\nlet test_embedding_forward () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.embedding ~vocab_size:10 ~embed_dim:4 ~scale:false () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let indices = Nx.create Nx.int32 [| 3 |] [| 0l; 5l; 2l |] in\n let y = apply_out m vars ~training:false indices in\n equal ~msg:\"embedding output shape\" (list int) [ 3; 4 ]\n (Array.to_list (Nx.shape y))\n\nlet test_compose_embedding_linear () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let emb = Layer.embedding ~vocab_size:10 ~embed_dim:4 ~scale:false () in\n let proj = Layer.linear ~in_features:4 ~out_features:2 () in\n let m = Layer.compose emb proj in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let indices = Nx.create Nx.int32 [| 3 |] [| 0l; 5l; 2l |] in\n let y, _ = Layer.apply m vars ~training:false indices in\n equal ~msg:\"compose embedding+linear output shape\" (list int) [ 3; 2 ]\n (Array.to_list (Nx.shape y))\n\n(* Activations *)\n\nlet test_relu () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.relu () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.create Nx.float32 [| 4 |] [| -2.0; -0.5; 0.0; 3.0 |] in\n let y = apply_out m vars ~training:false x in\n let expected = Nx.create Nx.float32 [| 4 |] [| 0.0; 0.0; 0.0; 3.0 |] in\n equal ~msg:\"relu\" bool true (tensor_close ~eps:1e-6 ~expected ~actual:y)\n\nlet test_activation_no_params () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let activations =\n [\n Layer.relu ();\n Layer.gelu ();\n Layer.silu ();\n Layer.tanh ();\n Layer.sigmoid ();\n ]\n in\n let assert_no_params (m : (float, float) Layer.t) =\n let vars = Layer.init m ~dtype:Nx.float32 in\n match (Layer.params vars, Layer.state vars) with\n | Ptree.List [], Ptree.List [] -> ()\n | _ -> fail \"expected empty params and state\"\n in\n List.iter (fun m -> assert_no_params m) activations\n\n(* Dropout *)\n\nlet test_dropout_eval_identity () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.dropout ~rate:0.99 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 10 |] in\n let y = apply_out m vars ~training:false x in\n equal ~msg:\"dropout eval = identity\" bool true\n (tensor_close ~eps:1e-6 ~expected:x ~actual:y)\n\nlet test_dropout_training () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.dropout ~rate:0.5 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 10 |] in\n let y = apply_out m vars ~training:true x in\n equal ~msg:\"dropout training shape\" (list int) [ 10 ]\n (Array.to_list (Nx.shape y))\n\nlet test_dropout_rate_bounds () =\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Layer.dropout ~rate:(-0.1) ()));\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Layer.dropout ~rate:1.0 ()))\n\n(* Flatten *)\n\nlet test_flatten_forward () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.flatten () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 2; 3; 4 |] in\n let y = apply_out m vars ~training:false x in\n equal ~msg:\"flatten shape\" (list int) [ 2; 12 ] (Array.to_list (Nx.shape y))\n\n(* Sequential *)\n\nlet test_sequential_init_structure () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m =\n Layer.sequential\n [\n Layer.linear ~in_features:4 ~out_features:3 ();\n Layer.relu ();\n Layer.linear ~in_features:3 ~out_features:2 ();\n ]\n in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let param_items = Ptree.List.items_exn (Layer.params vars) in\n let state_items = Ptree.List.items_exn (Layer.state vars) in\n equal ~msg:\"sequential params length\" int 3 (List.length param_items);\n equal ~msg:\"sequential state length\" int 3 (List.length state_items);\n let f0 = Ptree.Dict.fields_exn (List.nth param_items 0) in\n let w0 = Ptree.Dict.get_tensor_exn f0 ~name:\"weight\" Nx.float32 in\n equal ~msg:\"layer0 weight shape\" (list int) [ 4; 3 ]\n (Array.to_list (Nx.shape w0));\n (match List.nth param_items 1 with\n | Ptree.List [] -> ()\n | _ -> fail \"relu should have no params\");\n let f2 = Ptree.Dict.fields_exn (List.nth param_items 2) in\n let w2 = Ptree.Dict.get_tensor_exn f2 ~name:\"weight\" Nx.float32 in\n equal ~msg:\"layer2 weight shape\" (list int) [ 3; 2 ]\n (Array.to_list (Nx.shape w2))\n\nlet test_sequential_forward () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m =\n Layer.sequential\n [\n Layer.linear ~in_features:4 ~out_features:3 ();\n Layer.relu ();\n Layer.linear ~in_features:3 ~out_features:2 ();\n ]\n in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 5; 4 |] in\n let y = apply_out m vars ~training:false x in\n equal ~msg:\"sequential output shape\" (list int) [ 5; 2 ]\n (Array.to_list (Nx.shape y))\n\n(* Convolution *)\n\nlet test_conv1d_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.conv1d ~in_channels:3 ~out_channels:8 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let w = Ptree.Dict.get_tensor_exn fields ~name:\"weight\" Nx.float32 in\n let b = Ptree.Dict.get_tensor_exn fields ~name:\"bias\" Nx.float32 in\n equal ~msg:\"weight shape\" (list int) [ 8; 3; 3 ] (Array.to_list (Nx.shape w));\n equal ~msg:\"bias shape\" (list int) [ 8 ] (Array.to_list (Nx.shape b))\n\nlet test_conv1d_forward () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.conv1d ~in_channels:2 ~out_channels:4 ~kernel_size:3 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 1; 2; 10 |] in\n let y = apply_out m vars ~training:false x in\n let shape = Nx.shape y in\n equal ~msg:\"conv1d output batch\" int 1 shape.(0);\n equal ~msg:\"conv1d output channels\" int 4 shape.(1);\n equal ~msg:\"conv1d output length\" int 10 shape.(2)\n\nlet test_conv2d_shapes () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.conv2d ~in_channels:3 ~out_channels:16 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let fields = Ptree.Dict.fields_exn (Layer.params vars) in\n let w = Ptree.Dict.get_tensor_exn fields ~name:\"weight\" Nx.float32 in\n let b = Ptree.Dict.get_tensor_exn fields ~name:\"bias\" Nx.float32 in\n equal ~msg:\"weight shape\" (list int) [ 16; 3; 3; 3 ]\n (Array.to_list (Nx.shape w));\n equal ~msg:\"bias shape\" (list int) [ 16 ] (Array.to_list (Nx.shape b))\n\nlet test_conv2d_forward () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.conv2d ~in_channels:1 ~out_channels:4 ~kernel_size:(3, 3) () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 1; 1; 8; 8 |] in\n let y = apply_out m vars ~training:false x in\n equal ~msg:\"conv2d output shape\" (list int) [ 1; 4; 8; 8 ]\n (Array.to_list (Nx.shape y))\n\n(* Pooling *)\n\nlet test_max_pool2d () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.max_pool2d ~kernel_size:(2, 2) () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 1; 1; 4; 4 |] in\n let y = apply_out m vars ~training:false x in\n equal ~msg:\"max_pool2d shape\" (list int) [ 1; 1; 2; 2 ]\n (Array.to_list (Nx.shape y))\n\nlet test_avg_pool2d () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.avg_pool2d ~kernel_size:(2, 2) () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 1; 1; 6; 6 |] in\n let y = apply_out m vars ~training:false x in\n equal ~msg:\"avg_pool2d shape\" (list int) [ 1; 1; 3; 3 ]\n (Array.to_list (Nx.shape y))\n\n(* Parameter count *)\n\nlet test_param_count () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let m = Layer.linear ~in_features:10 ~out_features:5 () in\n let vars = Layer.init m ~dtype:Nx.float32 in\n equal ~msg:\"linear param count\" int 55\n (Ptree.count_parameters (Layer.params vars))\n\nlet () =\n run \"Kaun.Layer\"\n [\n group \"linear\"\n [\n test \"shapes\" test_linear_shapes;\n test \"forward\" test_linear_forward;\n test \"manual params\" test_linear_manual_params;\n test \"param count\" test_param_count;\n ];\n group \"normalization\"\n [\n test \"layer_norm shapes\" test_layer_norm_shapes;\n test \"layer_norm forward\" test_layer_norm_forward;\n test \"rms_norm shapes\" test_rms_norm_shapes;\n test \"batch_norm shapes\" test_batch_norm_shapes;\n test \"batch_norm rank3 axes\" test_batch_norm_rank3_axes;\n test \"batch_norm running stats eval\"\n test_batch_norm_running_stats_eval;\n test \"batch_norm eval affine rank3\" test_batch_norm_eval_affine_rank3;\n ];\n group \"embedding\"\n [\n test \"shapes\" test_embedding_shapes;\n test \"forward\" test_embedding_forward;\n test \"compose embedding+linear\" test_compose_embedding_linear;\n ];\n group \"activation\"\n [ test \"relu\" test_relu; test \"no params\" test_activation_no_params ];\n group \"regularization\"\n [\n test \"dropout eval identity\" test_dropout_eval_identity;\n test \"dropout training\" test_dropout_training;\n test \"dropout rate bounds\" test_dropout_rate_bounds;\n ];\n group \"conv\"\n [\n test \"conv1d shapes\" test_conv1d_shapes;\n test \"conv1d forward\" test_conv1d_forward;\n test \"conv2d shapes\" test_conv2d_shapes;\n test \"conv2d forward\" test_conv2d_forward;\n ];\n group \"pooling\"\n [ test \"max_pool2d\" test_max_pool2d; test \"avg_pool2d\" test_avg_pool2d ];\n group \"reshape\" [ test \"flatten\" test_flatten_forward ];\n group \"sequential\"\n [\n test \"init structure\" test_sequential_init_structure;\n test \"forward\" test_sequential_forward;\n ];\n ]\n"
},
{
"path": "packages/kaun/test/test_loss.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Loss = Kaun.Loss\n\nlet flatten_f32 t = Nx.to_array (Nx.reshape [| -1 |] (Nx.cast Nx.float32 t))\n\nlet tensor_close ~eps ~expected ~actual =\n let xs = flatten_f32 expected in\n let ys = flatten_f32 actual in\n let nx = Array.length xs in\n let ny = Array.length ys in\n if nx <> ny then false\n else\n let ok = ref true in\n for i = 0 to nx - 1 do\n if abs_float (xs.(i) -. ys.(i)) > eps then ok := false\n done;\n !ok\n\nlet test_cross_entropy_known_value () =\n let logits = Nx.create Nx.float32 [| 1; 3 |] [| 2.0; 0.0; -2.0 |] in\n let labels = Nx.create Nx.float32 [| 1; 3 |] [| 1.0; 0.0; 0.0 |] in\n let expected = log (1.0 +. exp (-2.0) +. exp (-4.0)) in\n let actual = Loss.cross_entropy logits labels |> Nx.item [] in\n equal ~msg:\"cross_entropy known value\" (float 1e-6) expected actual\n\nlet test_sparse_matches_dense_2d () =\n let logits =\n Nx.create Nx.float32 [| 3; 4 |]\n [| 2.0; 0.1; -1.0; 0.3; 0.2; 1.7; -0.4; 0.9; -0.1; 0.8; 1.4; -2.0 |]\n in\n let indices = Nx.create Nx.int32 [| 3 |] [| 0l; 1l; 2l |] in\n let one_hot = Nx.cast Nx.float32 (Nx.one_hot ~num_classes:4 indices) in\n let dense = Loss.cross_entropy logits one_hot in\n let sparse = Loss.cross_entropy_sparse logits indices in\n equal ~msg:\"sparse = dense (2d)\" (float 1e-6) (Nx.item [] dense)\n (Nx.item [] sparse)\n\nlet test_sparse_matches_dense_nd () =\n let logits =\n Nx.create Nx.float32 [| 2; 2; 3 |]\n [| 0.7; -1.2; 0.5; 1.1; 0.2; -0.3; -0.8; 0.9; 0.4; 0.6; -0.5; 1.7 |]\n in\n let indices = Nx.create Nx.int32 [| 2; 2 |] [| 0l; 2l; 1l; 2l |] in\n let one_hot = Nx.cast Nx.float32 (Nx.one_hot ~num_classes:3 indices) in\n let dense = Loss.cross_entropy logits one_hot in\n let sparse = Loss.cross_entropy_sparse logits indices in\n equal ~msg:\"sparse = dense (nd)\" (float 1e-6) (Nx.item [] dense)\n (Nx.item [] sparse)\n\nlet test_cross_entropy_rejects_invalid_shapes () =\n raises_invalid_arg \"Loss.cross_entropy: logits must have rank >= 1\" (fun () ->\n let logits = Nx.scalar Nx.float32 0.0 in\n let labels = Nx.scalar Nx.float32 1.0 in\n ignore (Loss.cross_entropy logits labels));\n raises_invalid_arg\n \"Loss.cross_entropy: labels rank mismatch (got 1, expected 2)\" (fun () ->\n let logits = Nx.zeros Nx.float32 [| 2; 3 |] in\n let labels = Nx.zeros Nx.float32 [| 2 |] in\n ignore (Loss.cross_entropy logits labels));\n raises_invalid_arg\n \"Loss.cross_entropy: labels shape mismatch at axis 0 (got 4, expected 2)\"\n (fun () ->\n let logits = Nx.zeros Nx.float32 [| 2; 3 |] in\n let labels = Nx.zeros Nx.float32 [| 4; 3 |] in\n ignore (Loss.cross_entropy logits labels));\n raises_invalid_arg\n \"Loss.cross_entropy: logits class dimension must be positive (got 0)\"\n (fun () ->\n let logits = Nx.zeros Nx.float32 [| 2; 0 |] in\n let labels = Nx.zeros Nx.float32 [| 2; 0 |] in\n ignore (Loss.cross_entropy logits labels))\n\nlet test_sparse_rejects_non_integer_labels () =\n let logits = Nx.zeros Nx.float32 [| 2; 3 |] in\n let bad = Nx.zeros Nx.float32 [| 2 |] in\n let msg =\n Printf.sprintf \"Loss.cross_entropy_sparse: expected integer labels, got %s\"\n (Nx_core.Dtype.to_string Nx.float32)\n in\n raises_invalid_arg msg (fun () ->\n ignore (Loss.cross_entropy_sparse logits bad))\n\nlet test_sparse_rejects_shape_mismatch () =\n let logits_2d = Nx.zeros Nx.float32 [| 2; 3 |] in\n let bad_rank = Nx.zeros Nx.int32 [| 2; 1 |] in\n raises_invalid_arg\n \"Loss.cross_entropy_sparse: labels rank mismatch (got 2, expected 1)\"\n (fun () -> ignore (Loss.cross_entropy_sparse logits_2d bad_rank));\n let logits_3d = Nx.zeros Nx.float32 [| 2; 3; 4 |] in\n let bad_shape = Nx.zeros Nx.int32 [| 2; 5 |] in\n raises_invalid_arg\n \"Loss.cross_entropy_sparse: labels shape mismatch at axis 1 (got 5, \\\n expected 3)\" (fun () ->\n ignore (Loss.cross_entropy_sparse logits_3d bad_shape))\n\nlet test_sparse_rejects_invalid_logits_shape () =\n raises_invalid_arg \"Loss.cross_entropy_sparse: logits must have rank >= 1\"\n (fun () ->\n let logits = Nx.scalar Nx.float32 0.0 in\n let labels = Nx.scalar Nx.int32 0l in\n ignore (Loss.cross_entropy_sparse logits labels));\n raises_invalid_arg\n \"Loss.cross_entropy_sparse: logits class dimension must be positive (got 0)\"\n (fun () ->\n let logits = Nx.zeros Nx.float32 [| 2; 0 |] in\n let labels = Nx.zeros Nx.int32 [| 2 |] in\n ignore (Loss.cross_entropy_sparse logits labels))\n\nlet test_binary_cross_entropy_logits_stable () =\n let logits =\n Nx.create Nx.float32 [| 5 |] [| 1000.0; -1000.0; 0.0; 50.0; -50.0 |]\n in\n let labels = Nx.create Nx.float32 [| 5 |] [| 1.0; 0.0; 1.0; 0.0; 1.0 |] in\n let loss = Loss.binary_cross_entropy logits labels |> Nx.item [] in\n let xs = [| 1000.0; -1000.0; 0.0; 50.0; -50.0 |] in\n let ys = [| 1.0; 0.0; 1.0; 0.0; 1.0 |] in\n let expected = ref 0.0 in\n for i = 0 to Array.length xs - 1 do\n let x = xs.(i) in\n let y = ys.(i) in\n let v = max x 0.0 -. (x *. y) +. log1p (exp (-.abs_float x)) in\n expected := !expected +. v\n done;\n expected := !expected /. float_of_int (Array.length xs);\n equal ~msg:\"binary_cross_entropy stable value\" (float 1e-6) !expected loss;\n equal ~msg:\"binary_cross_entropy finite\" bool true\n (match classify_float loss with FP_nan | FP_infinite -> false | _ -> true)\n\nlet test_binary_cross_entropy_rejects_invalid_shapes () =\n raises_invalid_arg\n \"Loss.binary_cross_entropy: labels rank mismatch (got 1, expected 2)\"\n (fun () ->\n let logits = Nx.zeros Nx.float32 [| 2; 1 |] in\n let labels = Nx.zeros Nx.float32 [| 2 |] in\n ignore (Loss.binary_cross_entropy logits labels));\n raises_invalid_arg\n \"Loss.binary_cross_entropy: labels shape mismatch at axis 0 (got 3, \\\n expected 2)\" (fun () ->\n let logits = Nx.zeros Nx.float32 [| 2; 1 |] in\n let labels = Nx.zeros Nx.float32 [| 3; 1 |] in\n ignore (Loss.binary_cross_entropy logits labels))\n\nlet test_mse_gradient_exact () =\n let predictions = Nx.create Nx.float32 [| 2; 2 |] [| 0.5; -1.0; 2.0; 3.0 |] in\n let targets = Nx.create Nx.float32 [| 2; 2 |] [| 0.0; 1.0; 1.0; 2.0 |] in\n let grad = Rune.grad (fun p -> Loss.mse p targets) predictions in\n let expected =\n Nx.create Nx.float32 [| 2; 2 |]\n [|\n 2.0 *. (0.5 -. 0.0) /. 4.0;\n 2.0 *. (-1.0 -. 1.0) /. 4.0;\n 2.0 *. (2.0 -. 1.0) /. 4.0;\n 2.0 *. (3.0 -. 2.0) /. 4.0;\n |]\n in\n equal ~msg:\"mse grad exact\" bool true\n (tensor_close ~eps:1e-6 ~expected ~actual:grad)\n\nlet test_cross_entropy_sparse_dense_gradient_match () =\n let logits =\n Nx.create Nx.float32 [| 2; 3 |] [| 2.0; 1.0; 0.5; -1.0; 0.2; 0.0 |]\n in\n let indices = Nx.create Nx.int32 [| 2 |] [| 0l; 2l |] in\n let one_hot = Nx.cast Nx.float32 (Nx.one_hot ~num_classes:3 indices) in\n let dense_grad = Rune.grad (fun x -> Loss.cross_entropy x one_hot) logits in\n let sparse_grad =\n Rune.grad (fun x -> Loss.cross_entropy_sparse x indices) logits\n in\n equal ~msg:\"cross_entropy sparse grad = dense grad\" bool true\n (tensor_close ~eps:1e-6 ~expected:dense_grad ~actual:sparse_grad)\n\nlet test_regression_values () =\n let predictions = Nx.create Nx.float32 [| 3 |] [| 1.0; 4.0; 3.0 |] in\n let targets = Nx.create Nx.float32 [| 3 |] [| 1.0; 1.0; 2.0 |] in\n equal ~msg:\"mse value\" (float 1e-6) (10.0 /. 3.0)\n (Nx.item [] (Loss.mse predictions targets));\n equal ~msg:\"mae value\" (float 1e-6) (4.0 /. 3.0)\n (Nx.item [] (Loss.mae predictions targets))\n\nlet test_regression_broadcasting () =\n let predictions =\n Nx.create Nx.float32 [| 2; 3 |] [| 0.0; 1.0; 2.0; 3.0; 4.0; 5.0 |]\n in\n let targets = Nx.create Nx.float32 [| 1; 3 |] [| 1.0; 1.0; 1.0 |] in\n equal ~msg:\"mse broadcast\" (float 1e-6) (31.0 /. 6.0)\n (Nx.item [] (Loss.mse predictions targets));\n equal ~msg:\"mae broadcast\" (float 1e-6) (11.0 /. 6.0)\n (Nx.item [] (Loss.mae predictions targets))\n\nlet test_mae_gradient_exact () =\n let predictions = Nx.create Nx.float32 [| 2; 2 |] [| 2.0; -1.0; 4.0; 0.0 |] in\n let targets = Nx.create Nx.float32 [| 2; 2 |] [| 1.0; 1.0; 2.0; -3.0 |] in\n let grad = Rune.grad (fun p -> Loss.mae p targets) predictions in\n let expected =\n Nx.create Nx.float32 [| 2; 2 |] [| 0.25; -0.25; 0.25; 0.25 |]\n in\n equal ~msg:\"mae grad exact\" bool true\n (tensor_close ~eps:1e-6 ~expected ~actual:grad)\n\nlet () =\n run \"Kaun.Loss\"\n [\n group \"cross-entropy\"\n [\n test \"cross entropy known value\" test_cross_entropy_known_value;\n test \"cross entropy rejects invalid shapes\"\n test_cross_entropy_rejects_invalid_shapes;\n test \"sparse matches dense (2d)\" test_sparse_matches_dense_2d;\n test \"sparse matches dense (nd)\" test_sparse_matches_dense_nd;\n test \"sparse rejects non-integer labels\"\n test_sparse_rejects_non_integer_labels;\n test \"sparse rejects shape mismatch\"\n test_sparse_rejects_shape_mismatch;\n test \"sparse rejects invalid logits shape\"\n test_sparse_rejects_invalid_logits_shape;\n test \"sparse/dense gradients match\"\n test_cross_entropy_sparse_dense_gradient_match;\n ];\n group \"binary\"\n [\n test \"binary cross entropy logits stable\"\n test_binary_cross_entropy_logits_stable;\n test \"binary cross entropy rejects invalid shapes\"\n test_binary_cross_entropy_rejects_invalid_shapes;\n ];\n group \"regression\"\n [\n test \"mse value + mae value\" test_regression_values;\n test \"mse/mae broadcasting\" test_regression_broadcasting;\n test \"mse gradient exact\" test_mse_gradient_exact;\n test \"mae gradient exact\" test_mae_gradient_exact;\n ];\n ]\n"
},
{
"path": "packages/kaun/test/test_metric.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Metric = Kaun.Metric\n\n(* Tracker *)\n\nlet test_tracker_observe_and_mean () =\n let t = Metric.tracker () in\n Metric.observe t \"loss\" 1.0;\n Metric.observe t \"loss\" 3.0;\n equal ~msg:\"mean of two\" (float 1e-10) 2.0 (Metric.mean t \"loss\")\n\nlet test_tracker_count () =\n let t = Metric.tracker () in\n Metric.observe t \"acc\" 0.9;\n Metric.observe t \"acc\" 0.8;\n Metric.observe t \"acc\" 0.7;\n equal ~msg:\"count\" int 3 (Metric.count t \"acc\")\n\nlet test_tracker_not_found () =\n let t = Metric.tracker () in\n raises Not_found (fun () -> ignore (Metric.mean t \"missing\"));\n raises Not_found (fun () -> ignore (Metric.count t \"missing\"))\n\nlet test_tracker_reset () =\n let t = Metric.tracker () in\n Metric.observe t \"x\" 1.0;\n Metric.reset t;\n equal ~msg:\"empty after reset\"\n (list (pair string (float 1e-10)))\n [] (Metric.to_list t)\n\nlet test_tracker_to_list_sorted () =\n let t = Metric.tracker () in\n Metric.observe t \"loss\" 0.5;\n Metric.observe t \"accuracy\" 0.9;\n Metric.observe t \"lr\" 0.001;\n let names = List.map fst (Metric.to_list t) in\n equal ~msg:\"sorted by name\" (list string) [ \"accuracy\"; \"loss\"; \"lr\" ] names\n\nlet test_tracker_summary () =\n let t = Metric.tracker () in\n Metric.observe t \"loss\" 0.4;\n Metric.observe t \"accuracy\" 0.9;\n let s = Metric.summary t in\n (* sorted: accuracy before loss *)\n equal ~msg:\"summary format\" string \"accuracy: 0.9000 loss: 0.4000\" s\n\n(* Dataset evaluation *)\n\nlet test_eval_mean () =\n let data = Kaun.Data.of_array [| 2.0; 4.0; 6.0 |] in\n let result = Metric.eval Fun.id data in\n equal ~msg:\"eval mean\" (float 1e-10) 4.0 result\n\nlet test_eval_empty_raises () =\n let data = Kaun.Data.of_array [||] in\n raises_invalid_arg \"Metric.eval: empty dataset\" (fun () ->\n ignore (Metric.eval Fun.id data))\n\nlet test_eval_many () =\n let data = Kaun.Data.of_array [| 1.0; 3.0 |] in\n let result =\n Metric.eval_many\n (fun x -> [ (\"double\", x *. 2.0); (\"half\", x /. 2.0) ])\n data\n in\n equal ~msg:\"double\" (float 1e-10) 4.0 (List.assoc \"double\" result);\n equal ~msg:\"half\" (float 1e-10) 1.0 (List.assoc \"half\" result)\n\nlet test_eval_many_empty_raises () =\n let data = Kaun.Data.of_array [||] in\n raises_invalid_arg \"Metric.eval_many: empty dataset\" (fun () ->\n ignore (Metric.eval_many (fun x -> [ (\"v\", x) ]) data))\n\n(* Accuracy *)\n\nlet test_accuracy_multiclass () =\n (* logits: batch=4, classes=3 *)\n let predictions =\n Nx.create Nx.float32 [| 4; 3 |]\n [|\n (* predicted class 2 *) 0.1;\n 0.2;\n 0.7;\n (* predicted class 0 *) 0.9;\n 0.05;\n 0.05;\n (* predicted class 1 *) 0.1;\n 0.8;\n 0.1;\n (* predicted class 0 *) 0.6;\n 0.2;\n 0.2;\n |]\n in\n (* targets: class indices *)\n let targets = Nx.create Nx.int32 [| 4 |] [| 2l; 0l; 0l; 0l |] in\n (* correct: sample 0 (2=2), sample 1 (0=0), sample 3 (0=0) = 3/4 *)\n equal ~msg:\"multiclass accuracy\" (float 1e-6) 0.75\n (Metric.accuracy predictions targets)\n\nlet test_accuracy_binary () =\n let predictions = Nx.create Nx.float32 [| 4 |] [| 0.8; 0.3; 0.6; 0.1 |] in\n let targets = Nx.create Nx.int32 [| 4 |] [| 1l; 0l; 1l; 1l |] in\n (* predicted: 1, 0, 1, 0; targets: 1, 0, 1, 1 => 3/4 correct *)\n equal ~msg:\"binary accuracy\" (float 1e-6) 0.75\n (Metric.accuracy predictions targets)\n\nlet test_binary_accuracy_default_threshold () =\n let predictions = Nx.create Nx.float32 [| 4 |] [| 0.8; 0.3; 0.6; 0.1 |] in\n let targets = Nx.create Nx.float32 [| 4 |] [| 1.0; 0.0; 1.0; 1.0 |] in\n equal ~msg:\"binary_accuracy default\" (float 1e-6) 0.75\n (Metric.binary_accuracy predictions targets)\n\nlet test_binary_accuracy_custom_threshold () =\n let predictions = Nx.create Nx.float32 [| 4 |] [| 0.8; 0.3; 0.6; 0.1 |] in\n let targets = Nx.create Nx.float32 [| 4 |] [| 1.0; 1.0; 1.0; 0.0 |] in\n (* threshold=0.25: predicted 1, 1, 1, 0; targets: 1, 1, 1, 0 => 4/4 *)\n equal ~msg:\"binary_accuracy threshold=0.25\" (float 1e-6) 1.0\n (Metric.binary_accuracy ~threshold:0.25 predictions targets)\n\n(* Precision / Recall / F1 *)\n\n(* Test scenario: 3 classes, 6 samples. predictions (logits): argmax gives [0;\n 1; 0; 2; 1; 0] targets: [0; 1; 1; 2; 0; 0]\n\n Confusion per class: class 0: TP=2, FP=1, FN=1, support=3 class 1: TP=1,\n FP=1, FN=1, support=2 class 2: TP=1, FP=0, FN=0, support=1\n\n Per-class precision: [2/3; 1/2; 1/1] Per-class recall: [2/3; 1/2; 1/1]\n Per-class f1: [2/3; 1/2; 1/1] *)\n\nlet prf_predictions () =\n Nx.create Nx.float32 [| 6; 3 |]\n [|\n (* pred 0 *) 0.8;\n 0.1;\n 0.1;\n (* pred 1 *) 0.1;\n 0.7;\n 0.2;\n (* pred 0 *) 0.6;\n 0.3;\n 0.1;\n (* pred 2 *) 0.1;\n 0.2;\n 0.7;\n (* pred 1 *) 0.2;\n 0.6;\n 0.2;\n (* pred 0 *) 0.5;\n 0.3;\n 0.2;\n |]\n\nlet prf_targets () = Nx.create Nx.int32 [| 6 |] [| 0l; 1l; 1l; 2l; 0l; 0l |]\n\nlet test_precision_macro () =\n let predictions = prf_predictions () in\n let targets = prf_targets () in\n (* macro = mean(2/3, 1/2, 1/1) = (2/3 + 1/2 + 1) / 3 *)\n let expected = ((2.0 /. 3.0) +. (1.0 /. 2.0) +. 1.0) /. 3.0 in\n equal ~msg:\"precision macro\" (float 1e-6) expected\n (Metric.precision Macro predictions targets)\n\nlet test_precision_micro () =\n let predictions = prf_predictions () in\n let targets = prf_targets () in\n (* micro = sum(TP) / (sum(TP) + sum(FP)) = 4 / (4 + 2) = 2/3 *)\n equal ~msg:\"precision micro\" (float 1e-6) (4.0 /. 6.0)\n (Metric.precision Micro predictions targets)\n\nlet test_precision_weighted () =\n let predictions = prf_predictions () in\n let targets = prf_targets () in\n (* weighted = (3 * 2/3 + 2 * 1/2 + 1 * 1) / 6 = (2 + 1 + 1) / 6 = 2/3 *)\n let expected =\n ((3.0 *. 2.0 /. 3.0) +. (2.0 *. 1.0 /. 2.0) +. (1.0 *. 1.0)) /. 6.0\n in\n equal ~msg:\"precision weighted\" (float 1e-6) expected\n (Metric.precision Weighted predictions targets)\n\nlet test_recall_macro () =\n let predictions = prf_predictions () in\n let targets = prf_targets () in\n let expected = ((2.0 /. 3.0) +. (1.0 /. 2.0) +. 1.0) /. 3.0 in\n equal ~msg:\"recall macro\" (float 1e-6) expected\n (Metric.recall Macro predictions targets)\n\nlet test_recall_micro () =\n let predictions = prf_predictions () in\n let targets = prf_targets () in\n (* micro recall = sum(TP) / (sum(TP) + sum(FN)) = 4 / (4 + 2) = 2/3 *)\n equal ~msg:\"recall micro\" (float 1e-6) (4.0 /. 6.0)\n (Metric.recall Micro predictions targets)\n\nlet test_f1_macro () =\n let predictions = prf_predictions () in\n let targets = prf_targets () in\n (* per-class f1 = [2/3; 1/2; 1] *)\n let expected = ((2.0 /. 3.0) +. (1.0 /. 2.0) +. 1.0) /. 3.0 in\n equal ~msg:\"f1 macro\" (float 1e-6) expected\n (Metric.f1 Macro predictions targets)\n\nlet test_f1_micro () =\n let predictions = prf_predictions () in\n let targets = prf_targets () in\n (* micro f1 = 2*sum(TP) / (2*sum(TP) + sum(FP) + sum(FN)) = 2*4 / (2*4 + 2 +\n 2) = 8/12 = 2/3 *)\n equal ~msg:\"f1 micro\" (float 1e-6) (8.0 /. 12.0)\n (Metric.f1 Micro predictions targets)\n\nlet test_micro_equals_accuracy () =\n (* For multiclass single-label, micro P = micro R = micro F1 = accuracy *)\n let predictions = prf_predictions () in\n let targets = prf_targets () in\n let acc = Metric.accuracy predictions targets in\n equal ~msg:\"micro precision = accuracy\" (float 1e-6) acc\n (Metric.precision Micro predictions targets);\n equal ~msg:\"micro recall = accuracy\" (float 1e-6) acc\n (Metric.recall Micro predictions targets);\n equal ~msg:\"micro f1 = accuracy\" (float 1e-6) acc\n (Metric.f1 Micro predictions targets)\n\nlet test_precision_zero_predictions () =\n (* class 2 has no predictions: pred=[0,1,0], targets=[0,1,2] *)\n let predictions =\n Nx.create Nx.float32 [| 3; 3 |]\n [| 0.8; 0.2; 0.0; 0.1; 0.9; 0.0; 0.6; 0.4; 0.0 |]\n in\n let targets = Nx.create Nx.int32 [| 3 |] [| 0l; 1l; 2l |] in\n (* class 0: TP=1, FP=1 => P=1/2 class 1: TP=1, FP=0 => P=1 class 2: TP=0, FP=0\n => P=0.0 (zero-div) macro = (1/2 + 1 + 0) / 3 = 0.5 *)\n equal ~msg:\"precision with missing class\" (float 1e-6) 0.5\n (Metric.precision Macro predictions targets)\n\nlet test_binary_f1 () =\n (* 2-class problem *)\n let predictions =\n Nx.create Nx.float32 [| 4; 2 |]\n [|\n 0.9;\n 0.1;\n (* pred 0 *)\n 0.3;\n 0.7;\n (* pred 1 *)\n 0.4;\n 0.6;\n (* pred 1 *)\n 0.8;\n 0.2;\n (* pred 0 *)\n |]\n in\n let targets = Nx.create Nx.int32 [| 4 |] [| 0l; 1l; 0l; 0l |] in\n (* class 0: TP=2, FP=0, FN=1 => P=1.0, R=2/3, F1=2*1*(2/3)/(1+2/3)=4/5 *)\n (* class 1: TP=1, FP=1, FN=0 => P=1/2, R=1.0, F1=2*(1/2)*1/(1/2+1)=2/3 *)\n let expected_macro = ((4.0 /. 5.0) +. (2.0 /. 3.0)) /. 2.0 in\n equal ~msg:\"binary f1 macro\" (float 1e-6) expected_macro\n (Metric.f1 Macro predictions targets)\n\nlet () =\n run \"Kaun.Metric\"\n [\n group \"tracker\"\n [\n test \"observe and mean\" test_tracker_observe_and_mean;\n test \"count\" test_tracker_count;\n test \"not found raises\" test_tracker_not_found;\n test \"reset\" test_tracker_reset;\n test \"to_list sorted\" test_tracker_to_list_sorted;\n test \"summary\" test_tracker_summary;\n ];\n group \"eval\"\n [\n test \"eval mean\" test_eval_mean;\n test \"eval empty raises\" test_eval_empty_raises;\n test \"eval_many\" test_eval_many;\n test \"eval_many empty raises\" test_eval_many_empty_raises;\n ];\n group \"accuracy\"\n [\n test \"multiclass\" test_accuracy_multiclass;\n test \"binary\" test_accuracy_binary;\n test \"binary_accuracy default\" test_binary_accuracy_default_threshold;\n test \"binary_accuracy custom threshold\"\n test_binary_accuracy_custom_threshold;\n ];\n group \"precision/recall/f1\"\n [\n test \"precision macro\" test_precision_macro;\n test \"precision micro\" test_precision_micro;\n test \"precision weighted\" test_precision_weighted;\n test \"recall macro\" test_recall_macro;\n test \"recall micro\" test_recall_micro;\n test \"f1 macro\" test_f1_macro;\n test \"f1 micro\" test_f1_micro;\n test \"micro = accuracy\" test_micro_equals_accuracy;\n test \"precision zero predictions\" test_precision_zero_predictions;\n test \"binary f1\" test_binary_f1;\n ];\n ]\n"
},
{
"path": "packages/kaun/test/test_optim.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Optim = Kaun.Optim\nmodule Ptree = Kaun.Ptree\nmodule Grad = Kaun.Grad\n\n(* Schedules *)\n\nlet test_constant_schedule () =\n let s = Vega.Schedule.constant 0.01 in\n equal ~msg:\"step 1\" (float 1e-10) 0.01 (s 1);\n equal ~msg:\"step 100\" (float 1e-10) 0.01 (s 100);\n equal ~msg:\"step 0\" (float 1e-10) 0.01 (s 0)\n\nlet test_cosine_decay () =\n let s = Vega.Schedule.cosine_decay ~init_value:0.1 ~decay_steps:100 () in\n equal ~msg:\"step 0\" (float 1e-10) 0.1 (s 0);\n equal ~msg:\"step 100 (fully decayed)\" (float 1e-10) 0.0 (s 100);\n equal ~msg:\"step 200 (past decay)\" (float 1e-10) 0.0 (s 200);\n let mid = s 50 in\n equal ~msg:\"step 50 (midpoint)\" (float 1e-6) 0.05 mid\n\nlet test_cosine_decay_alpha () =\n let s =\n Vega.Schedule.cosine_decay ~init_value:0.1 ~decay_steps:100 ~alpha:0.1 ()\n in\n equal ~msg:\"step 100 (alpha floor)\" (float 1e-10) 0.01 (s 100)\n\nlet test_warmup_cosine () =\n let s =\n Vega.Schedule.warmup_cosine ~init_value:0.0 ~peak_value:0.01\n ~warmup_steps:100\n in\n equal ~msg:\"step 0\" (float 1e-10) 0.0 (s 0);\n equal ~msg:\"step 100 (peak)\" (float 1e-10) 0.01 (s 100);\n equal ~msg:\"step 200 (past warmup)\" (float 1e-10) 0.01 (s 200)\n\nlet test_warmup_linear () =\n let s = Vega.Schedule.linear ~init_value:0.0 ~end_value:0.1 ~steps:10 in\n equal ~msg:\"step 0\" (float 1e-10) 0.0 (s 0);\n equal ~msg:\"step 5 (midpoint)\" (float 1e-10) 0.05 (s 5);\n equal ~msg:\"step 10 (peak)\" (float 1e-10) 0.1 (s 10);\n equal ~msg:\"step 20 (past warmup)\" (float 1e-10) 0.1 (s 20)\n\nlet test_exponential_decay () =\n let s =\n Vega.Schedule.exponential_decay ~init_value:1.0 ~decay_rate:0.5\n ~decay_steps:10\n in\n equal ~msg:\"step 0\" (float 1e-10) 1.0 (s 0);\n equal ~msg:\"step 10\" (float 1e-6) 0.5 (s 10);\n equal ~msg:\"step 20\" (float 1e-6) 0.25 (s 20)\n\n(* Helpers *)\n\nlet quadratic_loss params =\n (* f(x) = 0.5 * sum(x^2), grad = x *)\n let (Ptree.P t) = Ptree.as_tensor_exn params in\n let t = Ptree.Tensor.to_typed_exn Nx.float32 (Ptree.P t) in\n Nx.mul (Nx.scalar Nx.float32 0.5) (Nx.sum (Nx.mul t t))\n\nlet make_params values =\n Ptree.tensor (Nx.create Nx.float32 [| Array.length values |] values)\n\nlet get_values params =\n let (Ptree.P t) = Ptree.as_tensor_exn params in\n let t = Ptree.Tensor.to_typed_exn Nx.float32 (Ptree.P t) in\n Nx.to_array (Nx.reshape [| -1 |] t)\n\nlet train_steps algo params ~steps =\n let state = Optim.init algo params in\n let p = ref params in\n let s = ref state in\n for _ = 1 to steps do\n let _loss, grads = Grad.value_and_grad quadratic_loss !p in\n let new_params, state' = Optim.step !s !p grads in\n p := new_params;\n s := state'\n done;\n !p\n\n(* SGD *)\n\nlet test_sgd_basic () =\n let lr = Vega.Schedule.constant 0.1 in\n let algo = Vega.sgd lr in\n let params = make_params [| 4.0; -3.0 |] in\n let result = train_steps algo params ~steps:1 in\n let v = get_values result in\n (* After 1 step: x - lr * x = x * (1 - lr) = x * 0.9 *)\n equal ~msg:\"sgd[0] after 1 step\" (float 1e-5) 3.6 v.(0);\n equal ~msg:\"sgd[1] after 1 step\" (float 1e-5) (-2.7) v.(1)\n\nlet test_sgd_converges () =\n let lr = Vega.Schedule.constant 0.1 in\n let algo = Vega.sgd lr in\n let params = make_params [| 10.0; -8.0 |] in\n let result = train_steps algo params ~steps:100 in\n let v = get_values result in\n equal ~msg:\"sgd converges[0]\" (float 1e-3) 0.0 v.(0);\n equal ~msg:\"sgd converges[1]\" (float 1e-3) 0.0 v.(1)\n\nlet test_sgd_momentum () =\n let lr = Vega.Schedule.constant 0.01 in\n let algo = Vega.sgd ~momentum:0.9 lr in\n let params = make_params [| 5.0; -3.0 |] in\n let result = train_steps algo params ~steps:100 in\n let v = get_values result in\n equal ~msg:\"sgd+momentum converges[0]\" (float 0.1) 0.0 v.(0);\n equal ~msg:\"sgd+momentum converges[1]\" (float 0.1) 0.0 v.(1)\n\nlet test_sgd_nesterov () =\n let lr = Vega.Schedule.constant 0.01 in\n let algo = Vega.sgd ~momentum:0.9 ~nesterov:true lr in\n let params = make_params [| 5.0; -3.0 |] in\n let result = train_steps algo params ~steps:100 in\n let v = get_values result in\n equal ~msg:\"sgd+nesterov converges[0]\" (float 1e-2) 0.0 v.(0);\n equal ~msg:\"sgd+nesterov converges[1]\" (float 1e-2) 0.0 v.(1)\n\n(* Adam *)\n\nlet test_adam_converges () =\n let lr = Vega.Schedule.constant 0.1 in\n let algo = Vega.adam lr in\n let params = make_params [| 5.0; -3.0 |] in\n let result = train_steps algo params ~steps:100 in\n let v = get_values result in\n equal ~msg:\"adam converges[0]\" (float 0.5) 0.0 v.(0);\n equal ~msg:\"adam converges[1]\" (float 0.5) 0.0 v.(1)\n\n(* AdamW *)\n\nlet test_adamw_converges () =\n let lr = Vega.Schedule.constant 0.1 in\n let algo = Vega.adamw lr in\n let params = make_params [| 5.0; -3.0 |] in\n let result = train_steps algo params ~steps:100 in\n let v = get_values result in\n equal ~msg:\"adamw converges[0]\" (float 0.5) 0.0 v.(0);\n equal ~msg:\"adamw converges[1]\" (float 0.5) 0.0 v.(1)\n\n(* RMSprop *)\n\nlet test_rmsprop_converges () =\n let lr = Vega.Schedule.constant 0.1 in\n let algo = Vega.rmsprop lr in\n let params = make_params [| 5.0; -3.0 |] in\n let result = train_steps algo params ~steps:100 in\n let v = get_values result in\n equal ~msg:\"rmsprop converges[0]\" (float 0.5) 0.0 v.(0);\n equal ~msg:\"rmsprop converges[1]\" (float 0.5) 0.0 v.(1)\n\nlet test_rmsprop_momentum () =\n let lr = Vega.Schedule.constant 0.01 in\n let algo = Vega.rmsprop ~momentum:0.9 lr in\n let params = make_params [| 5.0; -3.0 |] in\n let result = train_steps algo params ~steps:100 in\n let v = get_values result in\n equal ~msg:\"rmsprop+momentum converges[0]\" (float 0.5) 0.0 v.(0);\n equal ~msg:\"rmsprop+momentum converges[1]\" (float 0.5) 0.0 v.(1)\n\n(* Adagrad *)\n\nlet test_adagrad_converges () =\n let lr = Vega.Schedule.constant 0.5 in\n let algo = Vega.adagrad lr in\n let params = make_params [| 5.0; -3.0 |] in\n let result = train_steps algo params ~steps:100 in\n let v = get_values result in\n equal ~msg:\"adagrad converges[0]\" (float 0.5) 0.0 v.(0);\n equal ~msg:\"adagrad converges[1]\" (float 0.5) 0.0 v.(1)\n\nlet test_invalid_hyperparameters () =\n let lr = Vega.Schedule.constant 0.1 in\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Vega.sgd ~momentum:1.0 lr));\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Vega.adam ~eps:0.0 lr));\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Vega.adamw ~weight_decay:(-0.1) lr));\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Vega.rmsprop ~decay:1.0 lr));\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n (fun () -> ignore (Vega.adagrad ~eps:0.0 lr))\n\n(* Gradient utilities *)\n\nlet test_global_norm () =\n let t =\n Ptree.dict\n [\n (\"a\", Ptree.tensor (Nx.create Nx.float32 [| 2 |] [| 3.0; 4.0 |]));\n (\"b\", Ptree.tensor (Nx.create Nx.float32 [| 1 |] [| 0.0 |]));\n ]\n in\n (* sqrt(9 + 16 + 0) = 5 *)\n equal ~msg:\"global_norm\" (float 1e-5) 5.0 (Optim.global_norm t)\n\nlet test_clip_by_global_norm () =\n let t = Ptree.tensor (Nx.create Nx.float32 [| 2 |] [| 3.0; 4.0 |]) in\n (* norm = 5, clip to 2.5 → scale by 0.5 *)\n let clipped = Optim.clip_by_global_norm 2.5 t in\n let v = get_values clipped in\n equal ~msg:\"clipped[0]\" (float 1e-5) 1.5 v.(0);\n equal ~msg:\"clipped[1]\" (float 1e-5) 2.0 v.(1)\n\nlet test_clip_no_op () =\n let t = Ptree.tensor (Nx.create Nx.float32 [| 2 |] [| 1.0; 1.0 |]) in\n (* norm = sqrt(2) ~ 1.41, max_norm = 5.0 → no clipping *)\n let clipped = Optim.clip_by_global_norm 5.0 t in\n let v = get_values clipped in\n equal ~msg:\"no clip[0]\" (float 1e-5) 1.0 v.(0);\n equal ~msg:\"no clip[1]\" (float 1e-5) 1.0 v.(1)\n\n(* Multi-parameter tree *)\n\nlet test_multi_param_tree () =\n let lr = Vega.Schedule.constant 0.1 in\n let algo = Vega.sgd lr in\n let params =\n Ptree.dict\n [\n (\"w\", Ptree.tensor (Nx.create Nx.float32 [| 2 |] [| 4.0; -2.0 |]));\n (\"b\", Ptree.tensor (Nx.create Nx.float32 [| 1 |] [| 1.0 |]));\n ]\n in\n let f p =\n let fields = Ptree.Dict.fields_exn p in\n let w = Ptree.Dict.get_tensor_exn fields ~name:\"w\" Nx.float32 in\n let b = Ptree.Dict.get_tensor_exn fields ~name:\"b\" Nx.float32 in\n Nx.add\n (Nx.mul (Nx.scalar Nx.float32 0.5) (Nx.sum (Nx.mul w w)))\n (Nx.mul (Nx.scalar Nx.float32 0.5) (Nx.sum (Nx.mul b b)))\n in\n let state = Optim.init algo params in\n let _loss, grads = Grad.value_and_grad f params in\n let result, _state' = Optim.step state params grads in\n let fields = Ptree.Dict.fields_exn result in\n let w = Ptree.Dict.get_tensor_exn fields ~name:\"w\" Nx.float32 in\n let b = Ptree.Dict.get_tensor_exn fields ~name:\"b\" Nx.float32 in\n (* w_new = w - lr * w = w * 0.9 *)\n equal ~msg:\"w[0]\" (float 1e-5) 3.6 (Nx.item [ 0 ] w);\n equal ~msg:\"w[1]\" (float 1e-5) (-1.8) (Nx.item [ 1 ] w);\n equal ~msg:\"b[0]\" (float 1e-5) 0.9 (Nx.item [ 0 ] b)\n\nlet () =\n run \"Kaun.Optim\"\n [\n group \"schedules\"\n [\n test \"constant\" test_constant_schedule;\n test \"cosine decay\" test_cosine_decay;\n test \"cosine decay alpha\" test_cosine_decay_alpha;\n test \"warmup cosine\" test_warmup_cosine;\n test \"warmup linear\" test_warmup_linear;\n test \"exponential decay\" test_exponential_decay;\n ];\n group \"sgd\"\n [\n test \"basic step\" test_sgd_basic;\n test \"converges\" test_sgd_converges;\n test \"momentum\" test_sgd_momentum;\n test \"nesterov\" test_sgd_nesterov;\n ];\n group \"adam\" [ test \"converges\" test_adam_converges ];\n group \"adamw\" [ test \"converges\" test_adamw_converges ];\n group \"rmsprop\"\n [\n test \"converges\" test_rmsprop_converges;\n test \"momentum\" test_rmsprop_momentum;\n ];\n group \"adagrad\" [ test \"converges\" test_adagrad_converges ];\n group \"validation\"\n [ test \"invalid hyperparameters\" test_invalid_hyperparameters ];\n group \"gradient utilities\"\n [\n test \"global_norm\" test_global_norm;\n test \"clip_by_global_norm\" test_clip_by_global_norm;\n test \"clip no-op\" test_clip_no_op;\n ];\n group \"multi-param\" [ test \"tree optimizer step\" test_multi_param_tree ];\n ]\n"
},
{
"path": "packages/kaun/test/test_ptree.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Ptree = Kaun.Ptree\n\nlet string_contains s sub =\n let slen = String.length s in\n let sub_len = String.length sub in\n let rec loop i =\n if i + sub_len > slen then false\n else if String.sub s i sub_len = sub then true\n else loop (i + 1)\n in\n if sub_len = 0 then true else loop 0\n\nlet raises_invalid_arg_any f =\n raises_match\n (fun exn -> match exn with Invalid_argument _ -> true | _ -> false)\n f\n\nlet raises_invalid_arg_contains needle f =\n raises_match\n (fun exn ->\n match exn with\n | Invalid_argument msg -> string_contains msg needle\n | _ -> false)\n f\n\nlet f32_leaf v = Ptree.tensor (Nx.full Nx.float32 [| 1 |] v)\n\nlet f32_value_of_tensor p =\n let t = Ptree.Tensor.to_typed_exn Nx.float32 p in\n Nx.item [ 0 ] t\n\nlet f32_value_of_tree t =\n let p = Ptree.as_tensor_exn t in\n f32_value_of_tensor p\n\nlet collect_f32_values t =\n let values = ref [] in\n Ptree.iter\n (fun p ->\n let v =\n Ptree.with_tensor p\n {\n run =\n (fun (type a) (type layout) (x : (a, layout) Nx.t) ->\n let y = Nx.cast Nx.float32 x in\n Nx.item [ 0 ] (Nx.reshape [| -1 |] y));\n }\n in\n values := v :: !values)\n t;\n List.rev !values\n\nlet test_dict_key_validation () =\n raises_invalid_arg \"duplicate key \\\"w\\\"\" (fun () ->\n Ptree.dict [ (\"w\", f32_leaf 1.0); (\"w\", f32_leaf 2.0) ]);\n raises_invalid_arg \"empty key\" (fun () -> Ptree.dict [ (\"\", f32_leaf 1.0) ]);\n raises_invalid_arg_contains \"reserved character '.'\" (fun () ->\n Ptree.dict [ (\"a.b\", f32_leaf 1.0) ]);\n raises_invalid_arg_contains \"reserved character '['\" (fun () ->\n Ptree.dict [ (\"a[0]\", f32_leaf 1.0) ]);\n ignore\n (Ptree.dict\n [\n (\"weight\", f32_leaf 1.0);\n (\"bias\", f32_leaf 2.0);\n (\"layer_1\", f32_leaf 3.0);\n ])\n\nlet test_tensor_module () =\n let p = Ptree.P (Nx.zeros Nx.float32 [| 2; 3 |]) in\n let dtype_matches =\n match Ptree.Tensor.dtype p with\n | Nx_core.Dtype.Pack dt -> Nx_core.Dtype.equal dt Nx.float32\n in\n equal ~msg:\"dtype\" bool true dtype_matches;\n equal ~msg:\"shape\" (list int) [ 2; 3 ] (Array.to_list (Ptree.Tensor.shape p));\n equal ~msg:\"numel\" int 6 (Ptree.Tensor.numel p);\n equal ~msg:\"to_typed hit\" bool true\n (Option.is_some (Ptree.Tensor.to_typed Nx.float32 p));\n equal ~msg:\"to_typed miss\" bool true\n (Option.is_none (Ptree.Tensor.to_typed Nx.float64 p));\n raises_invalid_arg_contains \"dtype mismatch\" (fun () ->\n Ptree.Tensor.to_typed_exn Nx.float64 p)\n\nlet test_leaf_access () =\n let p = Ptree.P (Nx.full Nx.float32 [| 1 |] 7.0) in\n let v =\n Ptree.with_tensor p\n {\n run =\n (fun (type a) (type layout) (t : (a, layout) Nx.t) ->\n let t = Nx.cast Nx.float32 t in\n Nx.item [ 0 ] t);\n }\n in\n equal ~msg:\"with_tensor\" (float 1e-6) 7.0 v;\n equal ~msg:\"as_tensor_exn\" (float 1e-6) 7.0\n (f32_value_of_tree (Ptree.tensor (Nx.full Nx.float32 [| 1 |] 7.0)));\n raises_invalid_arg_contains \"ctx\" (fun () ->\n Ptree.as_tensor_exn ~ctx:\"ctx\" (Ptree.list []))\n\nlet test_dict_access () =\n let fields =\n Ptree.Dict.fields_exn\n (Ptree.dict [ (\"w\", f32_leaf 3.0); (\"b\", f32_leaf 4.0) ])\n in\n equal ~msg:\"find hit\" bool true (Option.is_some (Ptree.Dict.find \"w\" fields));\n equal ~msg:\"find miss\" bool true (Option.is_none (Ptree.Dict.find \"x\" fields));\n equal ~msg:\"find_exn\" (float 1e-6) 4.0\n (f32_value_of_tree (Ptree.Dict.find_exn \"b\" fields));\n raises_invalid_arg_contains \"ctx\" (fun () ->\n Ptree.Dict.find_exn ~ctx:\"ctx\" \"x\" fields);\n equal ~msg:\"get_tensor_exn\" (float 1e-6) 3.0\n (Nx.item [ 0 ] (Ptree.Dict.get_tensor_exn fields ~name:\"w\" Nx.float32));\n raises_invalid_arg_any (fun () ->\n Ptree.Dict.get_tensor_exn fields ~name:\"x\" Nx.float32);\n raises_invalid_arg_any (fun () ->\n Ptree.Dict.get_tensor_exn fields ~name:\"w\" Nx.float64);\n raises_invalid_arg_any (fun () ->\n Ptree.Dict.get_tensor_exn\n (Ptree.Dict.fields_exn (Ptree.dict [ (\"node\", Ptree.list []) ]))\n ~name:\"node\" Nx.float32);\n raises_invalid_arg_contains \"ctx\" (fun () ->\n Ptree.Dict.fields_exn ~ctx:\"ctx\"\n (Ptree.tensor (Nx.zeros Nx.float32 [| 1 |])))\n\nlet test_list_access () =\n let items =\n Ptree.List.items_exn\n (Ptree.list [ f32_leaf 1.0; f32_leaf 2.0; Ptree.list [ f32_leaf 3.0 ] ])\n in\n equal ~msg:\"items_exn length\" int 3 (List.length items);\n raises_invalid_arg_contains \"ctx\" (fun () ->\n Ptree.List.items_exn ~ctx:\"ctx\" (f32_leaf 1.0))\n\nlet test_map () =\n let tree =\n Ptree.dict\n [ (\"a\", f32_leaf 1.0); (\"b\", Ptree.list [ f32_leaf 2.0; f32_leaf 3.0 ]) ]\n in\n let mapped =\n Ptree.map\n {\n run =\n (fun (type a) (type layout) (t : (a, layout) Nx.t) ->\n let dt = Nx.dtype t in\n let ten = Nx_core.Dtype.of_float dt 10.0 in\n Nx.add t (Nx.scalar dt ten));\n }\n tree\n in\n equal ~msg:\"map values\"\n (list (float 1e-6))\n [ 11.0; 12.0; 13.0 ]\n (collect_f32_values mapped)\n\nlet test_map2_success_and_order () =\n let lhs = Ptree.dict [ (\"z\", f32_leaf 1.0); (\"a\", f32_leaf 2.0) ] in\n let rhs = Ptree.dict [ (\"a\", f32_leaf 20.0); (\"z\", f32_leaf 10.0) ] in\n let out = Ptree.map2 { run = Nx.add } lhs rhs in\n let fields = Ptree.Dict.fields_exn out in\n equal ~msg:\"preserve lhs key order\" (list string) [ \"z\"; \"a\" ]\n (List.map fst fields);\n equal ~msg:\"z value\" (float 1e-6) 11.0\n (Nx.item [ 0 ] (Ptree.Dict.get_tensor_exn fields ~name:\"z\" Nx.float32));\n equal ~msg:\"a value\" (float 1e-6) 22.0\n (Nx.item [ 0 ] (Ptree.Dict.get_tensor_exn fields ~name:\"a\" Nx.float32))\n\nlet test_map2_errors () =\n raises_invalid_arg_contains \"structure mismatch\" (fun () ->\n Ptree.map2 { run = Nx.add } (f32_leaf 1.0)\n (Ptree.dict [ (\"x\", f32_leaf 1.0) ]));\n raises_invalid_arg_contains \"list length mismatch\" (fun () ->\n Ptree.map2 { run = Nx.add }\n (Ptree.list [ f32_leaf 1.0 ])\n (Ptree.list [ f32_leaf 1.0; f32_leaf 2.0 ]));\n raises_invalid_arg_contains \"dict size mismatch\" (fun () ->\n Ptree.map2 { run = Nx.add }\n (Ptree.dict [ (\"a\", f32_leaf 1.0) ])\n (Ptree.dict [ (\"a\", f32_leaf 1.0); (\"b\", f32_leaf 2.0) ]));\n raises_invalid_arg_contains \"not found in second dict\" (fun () ->\n Ptree.map2 { run = Nx.add }\n (Ptree.dict [ (\"a\", f32_leaf 1.0) ])\n (Ptree.dict [ (\"b\", f32_leaf 1.0) ]));\n raises_invalid_arg_contains \"dtype mismatch\" (fun () ->\n Ptree.map2 { run = Nx.add }\n (Ptree.tensor (Nx.ones Nx.float32 [| 1 |]))\n (Ptree.tensor (Nx.ones Nx.int32 [| 1 |])))\n\nlet test_iter_and_fold_order () =\n let tree =\n Ptree.dict\n [\n (\"a\", f32_leaf 1.0);\n (\"b\", Ptree.list [ f32_leaf 2.0; Ptree.dict [ (\"c\", f32_leaf 3.0) ] ]);\n (\"d\", f32_leaf 4.0);\n ]\n in\n let iter_values = collect_f32_values tree in\n equal ~msg:\"iter order\" (list (float 1e-6)) [ 1.0; 2.0; 3.0; 4.0 ] iter_values;\n let fold_values =\n Ptree.fold\n (fun acc p ->\n let v = f32_value_of_tensor p in\n v :: acc)\n [] tree\n |> List.rev\n in\n equal ~msg:\"fold order\" (list (float 1e-6)) [ 1.0; 2.0; 3.0; 4.0 ] fold_values\n\nlet test_flatten_and_rebuild () =\n let tree =\n Ptree.dict\n [ (\"a\", f32_leaf 1.0); (\"b\", Ptree.list [ f32_leaf 2.0; f32_leaf 3.0 ]) ]\n in\n let leaves, rebuild = Ptree.flatten tree in\n equal ~msg:\"flatten length\" int 3 (List.length leaves);\n equal ~msg:\"flatten order\"\n (list (float 1e-6))\n [ 1.0; 2.0; 3.0 ]\n (List.map f32_value_of_tensor leaves);\n let leaves_plus_10 =\n List.map\n (fun (Ptree.P t) ->\n Ptree.P\n (Nx.add t\n (Nx.scalar (Nx.dtype t) (Nx_core.Dtype.of_float (Nx.dtype t) 10.0))))\n leaves\n in\n let rebuilt = rebuild leaves_plus_10 in\n equal ~msg:\"rebuild values\"\n (list (float 1e-6))\n [ 11.0; 12.0; 13.0 ]\n (collect_f32_values rebuilt);\n let first_leaf =\n match leaves with\n | first :: _ -> first\n | [] -> fail \"flatten returned no leaves\"\n in\n raises_invalid_arg_contains \"not enough tensors\" (fun () ->\n rebuild [ first_leaf ]);\n raises_invalid_arg_contains \"too many tensors\" (fun () ->\n rebuild (leaves @ [ first_leaf ]))\n\nlet test_flatten_with_paths () =\n let root = f32_leaf 3.0 in\n equal ~msg:\"tensor root path\" (list string) [ \"\" ]\n (List.map fst (Ptree.flatten_with_paths root));\n let tree =\n Ptree.dict\n [\n (\"w\", f32_leaf 1.0);\n ( \"layers\",\n Ptree.list [ f32_leaf 2.0; Ptree.dict [ (\"b\", f32_leaf 3.0) ] ] );\n ]\n in\n equal ~msg:\"nested paths\" (list string)\n [ \"w\"; \"layers.0\"; \"layers.1.b\" ]\n (List.map fst (Ptree.flatten_with_paths tree))\n\nlet test_zeros_like_and_count_parameters () =\n let tree =\n Ptree.dict\n [\n (\"w\", Ptree.tensor (Nx.ones Nx.float32 [| 2; 3 |]));\n (\"b\", Ptree.tensor (Nx.full Nx.float32 [| 4 |] 5.0));\n ]\n in\n equal ~msg:\"count parameters\" int 10 (Ptree.count_parameters tree);\n let zeros = Ptree.zeros_like tree in\n equal ~msg:\"count preserved\" int 10 (Ptree.count_parameters zeros);\n Ptree.iter\n (fun p ->\n Ptree.with_tensor p\n {\n run =\n (fun (type a) (type layout) (x : (a, layout) Nx.t) ->\n let y = Nx.cast Nx.float32 x in\n let flat = Nx.reshape [| -1 |] y in\n let n = Nx.numel flat in\n for i = 0 to n - 1 do\n equal ~msg:\"zeros_like values\" (float 1e-6) 0.0\n (Nx.item [ i ] flat)\n done);\n })\n zeros\n\nlet test_pp () =\n let tree =\n Ptree.dict\n [\n (\"w\", Ptree.tensor (Nx.ones Nx.float32 [| 2 |]));\n (\"b\", Ptree.list [ f32_leaf 1.0 ]);\n ]\n in\n let s = Format.asprintf \"%a\" Ptree.pp tree in\n equal ~msg:\"pp has Dict\" bool true\n (String.length s >= String.length \"Dict\" && String.sub s 0 4 = \"Dict\");\n equal ~msg:\"pp has key\" bool true (string_contains s \"w\");\n equal ~msg:\"pp non-empty\" bool true (String.length s > 0)\n\nlet () =\n run \"Kaun.Ptree\"\n [\n group \"constructors\"\n [\n test \"dict key validation\" test_dict_key_validation;\n test \"leaf access\" test_leaf_access;\n ];\n group \"tensor\" [ test \"Tensor module\" test_tensor_module ];\n group \"containers\"\n [\n test \"dict access\" test_dict_access;\n test \"list access\" test_list_access;\n ];\n group \"functional\"\n [\n test \"map\" test_map;\n test \"map2 success and order\" test_map2_success_and_order;\n test \"map2 errors\" test_map2_errors;\n test \"iter/fold traversal order\" test_iter_and_fold_order;\n ];\n group \"flatten\"\n [\n test \"flatten/rebuild\" test_flatten_and_rebuild;\n test \"flatten_with_paths\" test_flatten_with_paths;\n ];\n group \"utilities\"\n [\n test \"zeros_like/count_parameters\"\n test_zeros_like_and_count_parameters;\n test \"pp\" test_pp;\n ];\n ]\n"
},
{
"path": "packages/kaun/test/test_train.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\nmodule Data = Kaun.Data\nmodule Layer = Kaun.Layer\nmodule Train = Kaun.Train\nmodule Loss = Kaun.Loss\n\nlet test_make_init () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let model =\n Layer.sequential\n [\n Layer.linear ~in_features:2 ~out_features:4 ();\n Layer.tanh ();\n Layer.linear ~in_features:4 ~out_features:1 ();\n ]\n in\n let optimizer = Vega.adam (Vega.Schedule.constant 0.01) in\n let trainer = Train.make ~model ~optimizer in\n let st = Train.init trainer ~dtype:Nx.float32 in\n let vars = Train.vars st in\n let param_count = Kaun.Ptree.count_parameters (Layer.params vars) in\n equal ~msg:\"has parameters\" bool true (param_count > 0)\n\nlet test_step () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let model =\n Layer.sequential\n [\n Layer.linear ~in_features:2 ~out_features:4 ();\n Layer.tanh ();\n Layer.linear ~in_features:4 ~out_features:1 ();\n ]\n in\n let optimizer = Vega.adam (Vega.Schedule.constant 0.01) in\n let trainer = Train.make ~model ~optimizer in\n let st = Train.init trainer ~dtype:Nx.float32 in\n let x =\n Nx.create Nx.float32 [| 4; 2 |] [| 0.; 0.; 0.; 1.; 1.; 0.; 1.; 1. |]\n in\n let y = Nx.create Nx.float32 [| 4; 1 |] [| 0.; 1.; 1.; 0. |] in\n let loss_val, st' =\n Train.step trainer st ~training:true\n ~loss:(fun pred -> Loss.binary_cross_entropy pred y)\n x\n in\n let loss_f = Nx.item [] loss_val in\n equal ~msg:\"loss is finite\" bool true (Float.is_finite loss_f);\n let vars0 = Train.vars st in\n let vars1 = Train.vars st' in\n equal ~msg:\"params changed\" bool false\n (Layer.params vars0 == Layer.params vars1)\n\nlet test_fit () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let model =\n Layer.sequential\n [\n Layer.linear ~in_features:2 ~out_features:4 ();\n Layer.tanh ();\n Layer.linear ~in_features:4 ~out_features:1 ();\n ]\n in\n let optimizer = Vega.adam (Vega.Schedule.constant 0.05) in\n let trainer = Train.make ~model ~optimizer in\n let st = Train.init trainer ~dtype:Nx.float32 in\n let x =\n Nx.create Nx.float32 [| 4; 2 |] [| 0.; 0.; 0.; 1.; 1.; 0.; 1.; 1. |]\n in\n let y = Nx.create Nx.float32 [| 4; 1 |] [| 0.; 1.; 1.; 0. |] in\n let st' =\n Train.fit trainer st\n (Data.repeat 30 (x, fun pred -> Loss.binary_cross_entropy pred y))\n in\n let pred = Train.predict trainer st' x |> Nx.sigmoid in\n let p0 = Nx.item [ 0; 0 ] pred in\n let p1 = Nx.item [ 1; 0 ] pred in\n let p2 = Nx.item [ 2; 0 ] pred in\n let p3 = Nx.item [ 3; 0 ] pred in\n equal ~msg:\"[0,0] -> ~0\" bool true (p0 < 0.4);\n equal ~msg:\"[0,1] -> ~1\" bool true (p1 > 0.6);\n equal ~msg:\"[1,0] -> ~1\" bool true (p2 > 0.6);\n equal ~msg:\"[1,1] -> ~0\" bool true (p3 < 0.4)\n\nlet test_predict () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let model =\n Layer.sequential\n [ Layer.linear ~in_features:3 ~out_features:2 (); Layer.relu () ]\n in\n let optimizer = Vega.sgd (Vega.Schedule.constant 0.01) in\n let trainer = Train.make ~model ~optimizer in\n let st = Train.init trainer ~dtype:Nx.float32 in\n let x = Nx.ones Nx.float32 [| 2; 3 |] in\n let y = Train.predict trainer st x in\n equal ~msg:\"predict shape\" (list int) [ 2; 2 ] (Array.to_list (Nx.shape y))\n\nlet test_fit_with_reporting () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let model =\n Layer.sequential [ Layer.linear ~in_features:2 ~out_features:1 () ]\n in\n let optimizer = Vega.sgd (Vega.Schedule.constant 0.01) in\n let trainer = Train.make ~model ~optimizer in\n let st = Train.init trainer ~dtype:Nx.float32 in\n let x = Nx.create Nx.float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |] in\n let y = Nx.create Nx.float32 [| 2; 1 |] [| 1.; 0. |] in\n let report_count = ref 0 in\n let _st' =\n Train.fit trainer st\n ~report:(fun ~step ~loss:_ _st ->\n if step mod 3 = 0 then incr report_count)\n (Data.repeat 10 (x, fun pred -> Loss.binary_cross_entropy pred y))\n in\n equal ~msg:\"report called 3 times (steps 3,6,9)\" int 3 !report_count\n\nlet test_fit_early_stop () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let model =\n Layer.sequential [ Layer.linear ~in_features:2 ~out_features:1 () ]\n in\n let optimizer = Vega.sgd (Vega.Schedule.constant 0.01) in\n let trainer = Train.make ~model ~optimizer in\n let st = Train.init trainer ~dtype:Nx.float32 in\n let x = Nx.create Nx.float32 [| 2; 2 |] [| 1.; 2.; 3.; 4. |] in\n let y = Nx.create Nx.float32 [| 2; 1 |] [| 1.; 0. |] in\n let last_step = ref 0 in\n let _st' =\n Train.fit trainer st\n ~report:(fun ~step ~loss:_ _st ->\n last_step := step;\n if step >= 15 then raise_notrace Train.Early_stop)\n (Data.repeat 100 (x, fun pred -> Loss.binary_cross_entropy pred y))\n in\n equal ~msg:\"stopped at step 15\" int 15 !last_step\n\nlet test_batch_norm_state_threading () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let model =\n Layer.sequential\n [\n Layer.linear ~in_features:2 ~out_features:4 ();\n Layer.batch_norm ~num_features:4 ();\n Layer.relu ();\n Layer.linear ~in_features:4 ~out_features:1 ();\n ]\n in\n let optimizer = Vega.adam (Vega.Schedule.constant 0.01) in\n let trainer = Train.make ~model ~optimizer in\n let st0 = Train.init trainer ~dtype:Nx.float32 in\n let x =\n Nx.create Nx.float32 [| 4; 2 |] [| 0.; 0.; 0.; 1.; 1.; 0.; 1.; 1. |]\n in\n let y = Nx.create Nx.float32 [| 4; 1 |] [| 0.; 1.; 1.; 0. |] in\n let _, st1 =\n Train.step trainer st0 ~training:true\n ~loss:(fun pred -> Loss.binary_cross_entropy pred y)\n x\n in\n let state0 = Layer.state (Train.vars st0) in\n let state1 = Layer.state (Train.vars st1) in\n equal ~msg:\"batch_norm state changed after step\" bool false (state0 == state1)\n\nlet () =\n run \"Kaun.Train\"\n [\n group \"make/init\" [ test \"make and init\" test_make_init ];\n group \"step\" [ test \"single step\" test_step ];\n group \"fit\"\n [\n test \"xor convergence\" test_fit;\n test \"reporting\" test_fit_with_reporting;\n test \"early stop\" test_fit_early_stop;\n ];\n group \"predict\" [ test \"shapes\" test_predict ];\n group \"stateful\"\n [ test \"batch_norm state threading\" test_batch_norm_state_threading ];\n ]\n"
},
{
"path": "packages/munin/README.md",
"content": "# Munin\n\nLocal experiment tracking for [Raven](https://github.com/raven-ml/raven).\n\nTrack metrics, save artifacts, and compare runs on the local filesystem.\nComes with a terminal dashboard for live monitoring. Data is plain JSON,\nreadable with `jq`. No server, no accounts.\n\n## Quick Start\n\n```ocaml\nlet session =\n Munin.Session.start ~experiment:\"mnist\"\n ~params:[ (\"lr\", `Float 0.001); (\"epochs\", `Int 10) ]\n ()\nin\nfor step = 1 to 1000 do\n let loss = train_step () in\n Munin.Session.log_metric session ~step \"train/loss\" loss\ndone;\nMunin.Session.finish session ()\n```\n\n```sh\nmunin watch # live terminal dashboard\nmunin compare a b c # side-by-side params + summary\n```\n\n## Features\n\n- **Scalar metrics** with `define_metric` for auto-computed summaries (min, max, mean, last) and custom x-axes\n- **Media logging** -- images, tables, and files at specific steps\n- **Versioned artifacts** with content-addressed storage, aliases, and lineage tracking\n- **Terminal dashboard** with live metric charts, system resource panels, and EMA smoothing\n- **CLI** -- `runs`, `show`, `compare`, `metrics`, `watch`, `artifacts`, `delete`, `gc`\n- **System monitoring** -- opt-in CPU and memory tracking via background thread\n- **Run grouping** for hyperparameter sweeps, parent/child for nested runs\n- **Provenance** -- git commit, command line, hostname, environment captured automatically\n- **Plain JSON storage** -- append-only JSONL event logs, `jq`-friendly\n\n## Libraries\n\n| Library | Description |\n|---------|-------------|\n| `munin` | Core tracking: Session, Run, Store, Artifact |\n| `munin.tui` | Terminal dashboard (`munin watch`) |\n| `munin.sys` | Background system monitoring (CPU, memory) |\n\n## Examples\n\n- **01-basic** -- Minimal run with scalar metrics\n- **02-metrics** -- Metric definitions, auto-summaries, epoch tracking\n- **03-artifacts** -- Versioned checkpoints with aliases and lineage\n- **04-media** -- Logging images and tables\n- **05-parameter-sweep** -- Run grouping for hyperparameter search\n- **06-inspect** -- Reading runs programmatically\n- **07-system-monitor** -- Background CPU and memory tracking\n\n## Contributing\n\nSee the [Raven monorepo README](../../README.md) for guidelines.\n\n## License\n\nISC License. See [LICENSE](../../LICENSE) for details.\n"
},
{
"path": "packages/munin/bin/dune",
"content": "(executable\n (name main)\n (public_name munin)\n (package munin)\n (libraries munin munin.tui cmdliner jsont jsont.bytesrw))\n"
},
{
"path": "packages/munin/bin/main.ml",
"content": "let pp_kv pairs =\n List.iter\n (fun (key, value) ->\n Printf.printf \" %s: %s\\n\" key (Format.asprintf \"%a\" Value.pp value))\n pairs\n\nlet pp_string_list items = Printf.printf \" %s\\n\" (String.concat \", \" items)\n\nlet string_of_status = function\n | `running -> \"running\"\n | `finished -> \"finished\"\n | `failed -> \"failed\"\n | `killed -> \"killed\"\n\nlet string_of_kind = function\n | `dataset -> \"dataset\"\n | `model -> \"model\"\n | `checkpoint -> \"checkpoint\"\n | `file -> \"file\"\n | `dir -> \"dir\"\n | `other -> \"other\"\n\nlet string_of_payload = function `file -> \"file\" | `dir -> \"dir\"\n\nlet value_to_float = function\n | `Float f -> Some f\n | `Int n -> Some (Float.of_int n)\n | `String s -> Float.of_string_opt s\n | `Bool _ -> None\n\nlet sorted_unique_keys runs f =\n let seen = Hashtbl.create 16 in\n List.iter\n (fun run -> List.iter (fun (k, _) -> Hashtbl.replace seen k ()) (f run))\n runs;\n Hashtbl.to_seq_keys seen |> List.of_seq |> List.sort String.compare\n\nlet run_root root f =\n let store = Store.open_ ?root () in\n f store\n\nlet runs_cmd =\n let doc = \"List tracked runs\" in\n let experiment =\n let doc = \"Limit the listing to a single experiment.\" in\n Cmdliner.Arg.(\n value & opt (some string) None & info [ \"experiment\" ] ~docv:\"NAME\" ~doc)\n in\n Cmdliner.Cmd.v\n (Cmdliner.Cmd.info \"runs\" ~doc)\n Cmdliner.Term.(\n const (fun root experiment ->\n run_root root (fun store ->\n Store.list_runs store ?experiment ()\n |> List.iter (fun run ->\n Printf.printf \"%s\\t%s\\t%s\\t%s\\t%s\\t%s\\n\" (Run.id run)\n (Run.experiment_name run)\n (string_of_status (Run.status run))\n (Option.value (Run.parent_id run) ~default:\"-\")\n (Option.value (Run.name run) ~default:\"-\")\n (Option.value (Run.provenance run).git_commit ~default:\"-\"))))\n $ Cmdliner.Arg.(\n value & opt (some string) None & info [ \"root\" ] ~docv:\"DIR\")\n $ experiment)\n\nlet show_cmd =\n let doc = \"Show one run\" in\n let run_id =\n Cmdliner.Arg.(required & pos 0 (some string) None & info [] ~docv:\"RUN_ID\")\n in\n Cmdliner.Cmd.v\n (Cmdliner.Cmd.info \"show\" ~doc)\n Cmdliner.Term.(\n const (fun root run_id ->\n run_root root (fun store ->\n match Store.find_run store run_id with\n | None ->\n Printf.eprintf \"munin: run not found: %s\\n\" run_id;\n exit 1\n | Some run ->\n Printf.printf \"id: %s\\n\" (Run.id run);\n Printf.printf \"experiment: %s\\n\" (Run.experiment_name run);\n Printf.printf \"name: %s\\n\"\n (Option.value (Run.name run) ~default:\"-\");\n Printf.printf \"parent: %s\\n\"\n (Option.value (Run.parent_id run) ~default:\"-\");\n Printf.printf \"status: %s\\n\"\n (string_of_status (Run.status run));\n Printf.printf \"started_at: %.0f\\n\" (Run.started_at run);\n Option.iter\n (Printf.printf \"ended_at: %.0f\\n\")\n (Run.ended_at run);\n Printf.printf \"resumable: %b\\n\" (Run.resumable run);\n Printf.printf \"notes: %s\\n\"\n (Option.value (Run.notes run) ~default:\"-\");\n let prov = Run.provenance run in\n Printf.printf \"command: %s\\n\" (String.concat \" \" prov.command);\n Printf.printf \"cwd: %s\\n\" prov.cwd;\n Printf.printf \"hostname: %s\\n\"\n (Option.value prov.hostname ~default:\"-\");\n Printf.printf \"pid: %d\\n\" prov.pid;\n Printf.printf \"git_commit: %s\\n\"\n (Option.value prov.git_commit ~default:\"-\");\n Printf.printf \"git_dirty: %s\\n\"\n (match prov.git_dirty with\n | None -> \"-\"\n | Some true -> \"true\"\n | Some false -> \"false\");\n Printf.printf \"env:\\n\";\n List.iter\n (fun (key, value) -> Printf.printf \" %s=%s\\n\" key value)\n prov.env;\n Printf.printf \"tags:\\n\";\n List.iter (Printf.printf \" %s\\n\") (Run.tags run);\n Printf.printf \"params:\\n\";\n pp_kv (Run.params run);\n Printf.printf \"summary:\\n\";\n pp_kv (Run.summary run);\n Printf.printf \"metric_keys:\\n\";\n pp_string_list (Run.metric_keys run);\n Printf.printf \"latest_metrics:\\n\";\n List.iter\n (fun (key, (metric : Run.metric)) ->\n Printf.printf \" %s: step=%d value=%g\\n\" key metric.step\n metric.value)\n (Run.latest_metrics run);\n Printf.printf \"children:\\n\";\n List.iter\n (fun child -> Printf.printf \" %s\\n\" (Run.id child))\n (Run.children run);\n Printf.printf \"output_artifacts:\\n\";\n List.iter\n (fun artifact ->\n Printf.printf \" %s %s aliases=[%s] consumers=[%s]\\n\"\n (Artifact.name artifact)\n (Artifact.version artifact)\n (String.concat \",\" (Artifact.aliases artifact))\n (String.concat \",\" (Artifact.consumer_run_ids artifact)))\n (Run.output_artifacts run);\n Printf.printf \"input_artifacts:\\n\";\n List.iter\n (fun artifact ->\n Printf.printf \" %s %s producer=%s\\n\"\n (Artifact.name artifact)\n (Artifact.version artifact)\n (Option.value\n (Artifact.producer_run_id artifact)\n ~default:\"-\"))\n (Run.input_artifacts run)))\n $ Cmdliner.Arg.(\n value & opt (some string) None & info [ \"root\" ] ~docv:\"DIR\")\n $ run_id)\n\nlet artifacts_cmd =\n let doc = \"List stored artifacts\" in\n let name =\n let doc = \"Limit the listing to a single artifact name.\" in\n Cmdliner.Arg.(\n value & opt (some string) None & info [ \"name\" ] ~docv:\"NAME\" ~doc)\n in\n Cmdliner.Cmd.v\n (Cmdliner.Cmd.info \"artifacts\" ~doc)\n Cmdliner.Term.(\n const (fun root name ->\n run_root root (fun store ->\n Store.list_artifacts store ?name ()\n |> List.iter (fun artifact ->\n Printf.printf \"%s\\t%s\\t%s\\t%s\\t%d\\t%s\\t%s\\n\"\n (Artifact.name artifact)\n (Artifact.version artifact)\n (string_of_kind (Artifact.kind artifact))\n (string_of_payload (Artifact.payload artifact))\n (Artifact.size_bytes artifact)\n (Option.value\n (Artifact.producer_run_id artifact)\n ~default:\"-\")\n (String.concat \",\" (Artifact.consumer_run_ids artifact)))))\n $ Cmdliner.Arg.(\n value & opt (some string) None & info [ \"root\" ] ~docv:\"DIR\")\n $ name)\n\nlet watch_cmd =\n let doc = \"Launch the live experiment dashboard\" in\n let experiment =\n let doc = \"Limit to runs in a single experiment.\" in\n Cmdliner.Arg.(\n value & opt (some string) None & info [ \"experiment\" ] ~docv:\"NAME\" ~doc)\n in\n let runs =\n Cmdliner.Arg.(value & pos_all string [] & info [] ~docv:\"RUN_ID\")\n in\n Cmdliner.Cmd.v\n (Cmdliner.Cmd.info \"watch\" ~doc)\n Cmdliner.Term.(\n const (fun root experiment runs ->\n Munin_tui.run ?root ?experiment ~runs ())\n $ Cmdliner.Arg.(\n value & opt (some string) None & info [ \"root\" ] ~docv:\"DIR\")\n $ experiment $ runs)\n\nlet compare_cmd =\n let doc = \"Compare runs side by side\" in\n let root =\n Cmdliner.Arg.(value & opt (some string) None & info [ \"root\" ] ~docv:\"DIR\")\n in\n let run_ids =\n Cmdliner.Arg.(non_empty & pos_all string [] & info [] ~docv:\"RUN_ID\")\n in\n Cmdliner.Cmd.v\n (Cmdliner.Cmd.info \"compare\" ~doc)\n Cmdliner.Term.(\n const (fun root run_ids ->\n run_root root (fun store ->\n let runs =\n List.filter_map (fun id -> Store.find_run store id) run_ids\n in\n if runs = [] then (\n Printf.eprintf \"munin: no runs found\\n\";\n exit 1);\n let run_label run =\n Option.value (Run.name run) ~default:(Run.id run)\n in\n (* Header *)\n Printf.printf \"key\";\n List.iter (fun run -> Printf.printf \"\\t%s\" (run_label run)) runs;\n Printf.printf \"\\n\";\n (* Params *)\n let param_keys = sorted_unique_keys runs Run.params in\n List.iter\n (fun key ->\n Printf.printf \"%s\" key;\n List.iter\n (fun run ->\n let v =\n match List.assoc_opt key (Run.params run) with\n | Some v -> Format.asprintf \"%a\" Value.pp v\n | None -> \"-\"\n in\n Printf.printf \"\\t%s\" v)\n runs;\n Printf.printf \"\\n\")\n param_keys;\n (* Summaries *)\n let summary_keys = sorted_unique_keys runs Run.summary in\n (* Collect goals from metric_defs *)\n let goals = Hashtbl.create 8 in\n List.iter\n (fun run ->\n List.iter\n (fun (key, (def : Run.metric_def)) ->\n match def.goal with\n | Some g -> Hashtbl.replace goals key g\n | None -> ())\n (Run.metric_defs run))\n runs;\n List.iter\n (fun key ->\n Printf.printf \"%s\" key;\n let values =\n List.map\n (fun run ->\n Run.find_summary run key\n |> Fun.flip Option.bind value_to_float)\n runs\n in\n (* Find best index *)\n let best_idx =\n match Hashtbl.find_opt goals key with\n | None -> None\n | Some goal ->\n let compare =\n match goal with\n | `Minimize -> fun a b -> Float.compare a b\n | `Maximize -> fun a b -> Float.compare b a\n in\n let best = ref None in\n List.iteri\n (fun i v ->\n match (v, !best) with\n | Some v, None -> best := Some (i, v)\n | Some v, Some (_, bv) ->\n if compare v bv < 0 then best := Some (i, v)\n | None, _ -> ())\n values;\n Option.map fst !best\n in\n List.iteri\n (fun i _ ->\n let s =\n match List.nth values i with\n | Some v ->\n let s = Printf.sprintf \"%g\" v in\n if Some i = best_idx then s ^ \"*\" else s\n | None -> \"-\"\n in\n Printf.printf \"\\t%s\" s)\n runs;\n Printf.printf \"\\n\")\n summary_keys))\n $ root $ run_ids)\n\nlet metrics_cmd =\n let doc = \"Show metric history\" in\n let root =\n Cmdliner.Arg.(value & opt (some string) None & info [ \"root\" ] ~docv:\"DIR\")\n in\n let run_id =\n Cmdliner.Arg.(required & pos 0 (some string) None & info [] ~docv:\"RUN_ID\")\n in\n let key =\n let doc = \"Metric key to dump history for.\" in\n Cmdliner.Arg.(\n value & opt (some string) None & info [ \"key\" ] ~docv:\"KEY\" ~doc)\n in\n let format =\n let doc = \"Output format: tsv (default), csv, or json.\" in\n Cmdliner.Arg.(\n value\n & opt (enum [ (\"tsv\", `Tsv); (\"csv\", `Csv); (\"json\", `Json) ]) `Tsv\n & info [ \"format\" ] ~docv:\"FORMAT\" ~doc)\n in\n Cmdliner.Cmd.v\n (Cmdliner.Cmd.info \"metrics\" ~doc)\n Cmdliner.Term.(\n const (fun root run_id key format ->\n run_root root (fun store ->\n match Store.find_run store run_id with\n | None ->\n Printf.eprintf \"munin: run not found: %s\\n\" run_id;\n exit 1\n | Some run -> (\n match key with\n | None ->\n (* Listing mode *)\n Printf.printf \"key\\tlatest_value\\tlatest_step\\tcount\\n\";\n List.iter\n (fun (key, (m : Run.metric)) ->\n let count =\n List.length (Run.metric_history run key)\n in\n Printf.printf \"%s\\t%g\\t%d\\t%d\\n\" key m.value m.step\n count)\n (Run.latest_metrics run)\n | Some key -> (\n let history = Run.metric_history run key in\n match format with\n | `Tsv ->\n Printf.printf \"step\\ttimestamp\\tvalue\\n\";\n List.iter\n (fun (m : Run.metric) ->\n Printf.printf \"%d\\t%.6f\\t%g\\n\" m.step m.timestamp\n m.value)\n history\n | `Csv ->\n Printf.printf \"step,timestamp,value\\n\";\n List.iter\n (fun (m : Run.metric) ->\n Printf.printf \"%d,%.6f,%g\\n\" m.step m.timestamp\n m.value)\n history\n | `Json ->\n Printf.printf \"[\";\n List.iteri\n (fun i (m : Run.metric) ->\n if i > 0 then Printf.printf \",\";\n Printf.printf\n \"{\\\"step\\\":%d,\\\"timestamp\\\":%.6f,\\\"value\\\":%g}\"\n m.step m.timestamp m.value)\n history;\n Printf.printf \"]\\n\"))))\n $ root $ run_id $ key $ format)\n\nlet () =\n exit\n (Cmdliner.Cmd.eval\n (Cmdliner.Cmd.group\n (Cmdliner.Cmd.info \"munin\" ~doc:\"Local experiment tracking for Raven\")\n [\n runs_cmd;\n show_cmd;\n artifacts_cmd;\n watch_cmd;\n compare_cmd;\n metrics_cmd;\n ]))\n"
},
{
"path": "packages/munin/doc/01-getting-started.md",
"content": "# Getting Started\n\nThis guide covers installation, key concepts, and a complete first\nexample that tracks a run, inspects it via the CLI, and compares two\nruns.\n\n## Installation\n\n\n```bash\nopam install munin\n```\n\nOr build from source:\n\n\n```bash\ngit clone https://github.com/raven-ml/raven\ncd raven && dune build munin\n```\n\n## Key Concepts\n\n**Session.** A session is the write handle for a single run. All\nmutations go through append-only events -- no direct state editing.\n`Session.start` opens a session, `Session.finish` closes it.\n`Session.with_run` wraps both and handles exceptions.\n\n**Run.** A run is the persisted, read-only view of a tracked\nexperiment. It materializes its state by replaying the event log.\n`Run.params`, `Run.summary`, `Run.metric_history`, and other\naccessors expose the data.\n\n**Store.** A store is the root directory containing all experiments\nand artifacts. `Store.open_` creates or opens it.\n`Store.list_runs`, `Store.find_run`, and `Store.latest_run`\ndiscover runs across experiments.\n\n**Artifact.** An artifact is a versioned, content-addressed file or\ndirectory. Versions are auto-incremented (v1, v2, ...). Aliases like\n`\"latest\"` or `\"best\"` resolve to a specific version.\n`Session.log_artifact` produces one, `Session.use_artifact` records\nconsumption.\n\n**Value.** Parameters, summaries, and metadata use a simple scalar\ntype: `` [`Bool of bool | `Int of int | `Float of float | `String of string] ``.\n\n## Example: First Tracked Run\n\nThis example starts a session, logs hyperparameters and metrics,\nsaves an artifact, then reads everything back.\n\n\n```ocaml\nopen Munin\n\nlet () =\n let session =\n Session.start ~experiment:\"demo\" ~name:\"baseline\"\n ~params:[ (\"lr\", `Float 0.001); (\"hidden\", `Int 64) ]\n ()\n in\n (* Log metrics at each step. *)\n Session.define_metric session \"loss\" ~summary:`Min ~goal:`Minimize ();\n\n for step = 1 to 50 do\n let loss = 1.0 /. Float.of_int step in\n let acc = 1.0 -. loss in\n Session.log_metrics session ~step [ (\"loss\", loss); (\"accuracy\", acc) ]\n done;\n\n (* Write a summary value explicitly. *)\n Session.set_summary session [ (\"note\", `String \"first run\") ];\n\n Session.finish session ();\n Printf.printf \"run: %s\\n\" (Run.id (Session.run session))\n```\n\nAfter running, inspect from the terminal:\n\n\n```sh\n# List all runs.\nmunin runs\n\n# Show full details for a run.\nmunin show \n\n# Dump metric history as TSV.\nmunin metrics --key loss\n\n# Export as JSON.\nmunin metrics --key loss --format json\n```\n\n## Example: Comparing Two Runs\n\nRun the same experiment with different hyperparameters, then compare.\n\n\n```ocaml\nopen Munin\n\nlet train ~name ~lr =\n Session.with_run ~experiment:\"demo\" ~name\n ~params:[ (\"lr\", `Float lr) ]\n @@ fun session ->\n Session.define_metric session \"loss\" ~summary:`Min ~goal:`Minimize ();\n for step = 1 to 50 do\n let loss = (1.0 /. Float.of_int step) *. (1.0 /. lr) in\n Session.log_metric session ~step \"loss\" loss\n done\n\nlet () =\n train ~name:\"slow\" ~lr:0.01;\n train ~name:\"fast\" ~lr:0.1\n```\n\nCompare them side by side:\n\n\n```sh\nmunin compare \n```\n\nThe compare command prints a table with parameters and summary values.\nWhen a metric has a `goal` declared, the best value is marked with `*`.\n\n## Provenance\n\nEvery run automatically captures:\n\n- The command line (`Sys.argv`)\n- Working directory\n- Hostname and PID\n- Git commit hash and dirty status\n\nPass `~capture_env:[\"CUDA_VISIBLE_DEVICES\"; \"OMP_NUM_THREADS\"]` to\n`Session.start` to also record specific environment variables.\n\n## Store Location\n\nBy default, runs are stored in `$XDG_DATA_HOME/raven/munin`. Override\nwith the `RAVEN_TRACKING_DIR` environment variable, or pass `~root`\nto `Session.start` and `Store.open_`.\n\n## Next Steps\n\n- [Tracking Metrics](../02-tracking/) -- scalars, metric definitions, media, Kaun integration\n- [Artifacts](../03-artifacts/) -- versioned files, aliases, lineage, deduplication\n"
},
{
"path": "packages/munin/doc/02-tracking.md",
"content": "# Tracking Metrics\n\nThis page covers scalar metric logging, metric definitions with\nsummaries and goals, media logging, and integration with Kaun's\ntraining loop.\n\n## Logging Scalars\n\n`Session.log_metric` records a single scalar at a given step.\n`Session.log_metrics` records several atomically with the same\ntimestamp.\n\n\n```ocaml\nopen Munin\n\nlet () =\n Session.with_run ~experiment:\"tracking-demo\" @@ fun session ->\n for step = 1 to 100 do\n let loss = 1.0 /. Float.of_int step in\n let acc = 1.0 -. loss in\n Session.log_metrics session ~step [ (\"loss\", loss); (\"accuracy\", acc) ]\n done\n```\n\nEach call appends to an event log. The `step` is your x-axis counter\n(typically the global training step). A wall-clock timestamp is added\nautomatically; pass `~timestamp` to override it.\n\nRead metrics back through the `Run` module:\n\n\n```ocaml\nlet run = Session.run session in\nRun.metric_keys run (* [\"accuracy\"; \"loss\"] *)\nRun.latest_metrics run (* latest value per key *)\nRun.metric_history run \"loss\" (* full chronological history *)\n```\n\n## Defining Metrics\n\n`Session.define_metric` declares how a metric should be summarized,\ncompared, and plotted. Call it once per key, before or after logging\nvalues.\n\n\n```ocaml\nSession.define_metric session \"loss\"\n ~summary:`Min\n ~goal:`Minimize\n ();\n\nSession.define_metric session \"accuracy\"\n ~summary:`Max\n ~goal:`Maximize\n ();\n```\n\n### Summary Modes\n\nThe `~summary` parameter controls the auto-computed run summary value:\n\n| Mode | Summary value |\n|--------|---------------|\n| `` `Min `` | Minimum over all samples |\n| `` `Max `` | Maximum over all samples |\n| `` `Mean `` | Arithmetic mean of all samples |\n| `` `Last `` | Most recent sample (default) |\n| `` `None `` | No auto-summary |\n\nWhen the run is loaded, the summary is computed from the full metric\nhistory. You do not need to compute it yourself.\n\n### Explicit Summaries\n\n`Session.set_summary` writes explicit summary values that always take\nprecedence over auto-computed ones:\n\n\n```ocaml\nSession.set_summary session\n [ (\"best_loss\", `Float 0.023); (\"note\", `String \"converged early\") ]\n```\n\nUse this for values that are not simple aggregations of a metric\nhistory, or for non-float summaries.\n\n### Goal\n\nThe `~goal` parameter declares whether lower (`` `Minimize ``) or\nhigher (`` `Maximize ``) values are better. It is used by:\n\n- `munin compare` to mark the best value with `*`\n- `munin watch` TUI for \"best\" badges\n- `Run_monitor.best` to find the best observation\n\n### Step Metric\n\nThe `~step_metric` parameter specifies another metric as the x-axis:\n\n\n```ocaml\nSession.define_metric session \"val/accuracy\"\n ~summary:`Max ~goal:`Maximize ~step_metric:\"epoch\" ();\n```\n\nThis tells renderers to plot `val/accuracy` against the `epoch`\nmetric instead of the raw step counter.\n\n## Epoch Tracking\n\nEpochs are not a special concept -- log them as a regular metric and\nreference them with `~step_metric`:\n\n\n```ocaml\nSession.define_metric session \"train/loss\"\n ~summary:`Min ~goal:`Minimize ~step_metric:\"epoch\" ();\nSession.define_metric session \"val/accuracy\"\n ~summary:`Max ~goal:`Maximize ~step_metric:\"epoch\" ();\n\nfor epoch = 1 to 10 do\n let steps_per_epoch = 100 in\n for batch = 1 to steps_per_epoch do\n let step = ((epoch - 1) * steps_per_epoch) + batch in\n let loss = 1.0 /. Float.of_int step in\n Session.log_metrics session ~step\n [ (\"train/loss\", loss); (\"epoch\", Float.of_int epoch) ]\n done;\n let step = epoch * steps_per_epoch in\n Session.log_metric session ~step \"val/accuracy\" (Float.of_int epoch *. 0.1)\ndone\n```\n\n## Media Logging\n\n### Images and Files\n\n`Session.log_media` copies a file into the run's `media/` directory\nand records it in the event log. The `~kind` is metadata for\nrenderers.\n\n\n```ocaml\n(* Log an image at a specific step. *)\nSession.log_media session ~step:100 ~key:\"viz/confusion\"\n ~kind:`Image ~path:\"/tmp/confusion_matrix.png\";\n\n(* Log a text file. *)\nSession.log_media session ~step:1 ~key:\"config\"\n ~kind:`File ~path:\"config.yaml\"\n```\n\nKeys may contain `/` separators to organize media into a hierarchy.\nThe file is stored at `/media/_.`.\n\nRead media back:\n\n\n```ocaml\nlet run = Session.run session in\nRun.media_keys run (* [\"config\"; \"viz/confusion\"] *)\nRun.media_history run \"viz/confusion\" (* list of media_entry records *)\n```\n\n### Structured Tables\n\n`Session.log_table` stores a table as a JSON file. Useful for\nconfusion matrices, per-class metrics, or data samples.\n\n\n```ocaml\nSession.log_table session ~step:1 ~key:\"results/per_class\"\n ~columns:[ \"class\"; \"precision\"; \"recall\"; \"f1\" ]\n ~rows:[\n [ `String \"cat\"; `Float 0.92; `Float 0.88; `Float 0.90 ];\n [ `String \"dog\"; `Float 0.89; `Float 0.93; `Float 0.91 ];\n [ `String \"bird\"; `Float 0.95; `Float 0.91; `Float 0.93 ];\n ]\n```\n\n## Integration with Kaun\n\nMunin has no compile-time dependency on Kaun. Integration happens\nthrough `Train.fit`'s `~report` callback:\n\n\n```ocaml\nopen Kaun\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let session =\n Munin.Session.start ~experiment:\"mnist\" ~name:\"cnn-adam\"\n ~params:[\n (\"lr\", `Float 0.001);\n (\"batch_size\", `Int 64);\n (\"optimizer\", `String \"adam\");\n ]\n ()\n in\n Munin.Session.define_metric session \"train/loss\"\n ~summary:`Min ~goal:`Minimize ();\n Munin.Session.define_metric session \"val/accuracy\"\n ~summary:`Max ~goal:`Maximize ();\n\n let (x_train, y_train), (x_test, y_test) = Kaun_datasets.mnist () in\n let trainer =\n Train.make ~model ~optimizer:(Vega.adam (Vega.Schedule.constant 0.001))\n in\n let st = ref (Train.init trainer ~dtype:Nx.float32) in\n\n for epoch = 1 to 3 do\n let train_data =\n Data.prepare ~shuffle:true ~batch_size:64 (x_train, y_train)\n |> Data.map (fun (x, y) ->\n (x, fun logits -> Loss.cross_entropy_sparse logits y))\n in\n st :=\n Train.fit trainer !st\n ~report:(fun ~step ~loss _st ->\n Munin.Session.log_metrics session ~step\n [ (\"train/loss\", loss); (\"epoch\", Float.of_int epoch) ])\n train_data;\n\n (* Evaluate and log validation accuracy. *)\n let test_batches = Data.prepare ~batch_size:64 (x_test, y_test) in\n let acc =\n Metric.eval\n (fun (x, y) ->\n let logits = Train.predict trainer !st x in\n Metric.accuracy logits y)\n test_batches\n in\n Munin.Session.log_metric session ~step:(epoch * 937)\n \"val/accuracy\" acc\n done;\n\n Munin.Session.finish session ()\n```\n\n## System Monitoring\n\n`Munin_sys.start` spawns a background thread that samples CPU and\nmemory usage every 15 seconds (configurable via `~interval`):\n\n\n```ocaml\nlet sysmon = Munin_sys.start session () in\n(* ... training ... *)\nMunin_sys.stop sysmon\n```\n\nLogged metrics: `sys/cpu_user`, `sys/cpu_system`, `sys/mem_used_pct`,\n`sys/mem_used_gb`, `sys/proc_cpu_pct`, `sys/proc_mem_mb`.\n\n## Next Steps\n\n- [Artifacts](../03-artifacts/) -- versioned files, aliases, lineage, deduplication\n"
},
{
"path": "packages/munin/doc/03-artifacts.md",
"content": "# Artifacts\n\nArtifacts are versioned, content-addressed files or directories with\ncross-run lineage tracking. Use them for datasets, model checkpoints,\nand any other outputs you want to version alongside your runs.\n\n## Logging an Artifact\n\n`Session.log_artifact` copies a file or directory into the blob store,\nassigns a version, and records it as an output of the current run.\n\n\n```ocaml\nopen Munin\n\nlet () =\n Session.with_run ~experiment:\"pipeline\" ~name:\"prepare-data\"\n ~tags:[ \"data\" ]\n @@ fun session ->\n (* ... produce a dataset file ... *)\n let _artifact =\n Session.log_artifact session\n ~name:\"measurements\"\n ~kind:`dataset\n ~path:\"data/measurements.csv\"\n ~metadata:[ (\"rows\", `Int 10000); (\"format\", `String \"csv\") ]\n ~aliases:[ \"latest\" ]\n ()\n in\n ()\n```\n\nParameters:\n\n- **`~name`** -- logical name for the artifact (e.g. `\"measurements\"`, `\"mnist-cnn\"`).\n- **`~kind`** -- one of `` `checkpoint ``, `` `model ``, `` `dataset ``, `` `file ``, `` `dir ``, `` `other ``.\n- **`~path`** -- path to the file or directory to store.\n- **`~metadata`** -- optional key-value pairs attached to the version.\n- **`~aliases`** -- optional alias list (e.g. `[\"latest\"; \"best\"]`).\n\n## Artifact Kinds\n\nThe `kind` field is a semantic label. It does not affect storage; all\nartifacts are stored the same way.\n\n| Kind | Use for |\n|------|---------|\n| `` `checkpoint `` | Training checkpoints (model + optimizer state) |\n| `` `model `` | Final model weights |\n| `` `dataset `` | Datasets and data splits |\n| `` `file `` | Single files (configs, logs, reports) |\n| `` `dir `` | Directory trees |\n| `` `other `` | Anything else |\n\n## Versioning\n\nEach call to `log_artifact` with the same `~name` creates a new\nversion: `v1`, `v2`, `v3`, and so on. Versions are immutable once\ncreated.\n\n\n```ocaml\n(* First call creates v1. *)\nlet v1 =\n Session.log_artifact session ~name:\"model\" ~kind:`model\n ~path:\"model_epoch1.safetensors\" ()\nin\n(* Second call creates v2. *)\nlet v2 =\n Session.log_artifact session ~name:\"model\" ~kind:`model\n ~path:\"model_epoch2.safetensors\" ()\nin\nPrintf.printf \"%s %s\\n\" (Artifact.version v1) (Artifact.version v2)\n(* prints: v1 v2 *)\n```\n\n## Aliases\n\nAliases are mutable pointers to a specific version. Common aliases:\n\n- `\"latest\"` -- the most recent version\n- `\"best\"` -- the best-performing version\n\nPass `~aliases` when logging to attach them. When a new version gets\nthe same alias, it moves from the old version to the new one.\n\nResolve an alias through the store:\n\n\n```ocaml\nlet store = Store.open_ () in\nmatch Store.find_artifact store ~name:\"model\" ~version:\"latest\" with\n| Some artifact ->\n Printf.printf \"resolved to %s\\n\" (Artifact.version artifact)\n| None ->\n Printf.printf \"not found\\n\"\n```\n\n`Store.find_artifact` accepts both explicit versions (`\"v2\"`) and\naliases (`\"latest\"`).\n\n## Content-Addressed Deduplication\n\nArtifact payloads are stored in a blob directory keyed by their\nSHA-256 digest. If two versions have identical content, only one copy\nis stored on disk.\n\n\n```ocaml\n(* These share the same blob if the file content is identical. *)\nlet a = Session.log_artifact session ~name:\"config\" ~kind:`file\n ~path:\"config.yaml\" () in\nlet b = Session.log_artifact session ~name:\"config\" ~kind:`file\n ~path:\"config.yaml\" () in\nassert (Artifact.digest a = Artifact.digest b)\n```\n\n`Store.gc` removes blobs that are no longer referenced by any artifact\nversion.\n\n## Lineage\n\n### Producer\n\nWhen you call `Session.log_artifact`, the current run is automatically\nrecorded as the producer:\n\n\n```ocaml\nlet artifact =\n Session.log_artifact session ~name:\"features\" ~kind:`dataset\n ~path:\"features.csv\" ()\nin\nArtifact.producer_run_id artifact (* Some \"\" *)\n```\n\n### Consumer\n\n`Session.use_artifact` records the current run as a consumer of an\nexisting artifact:\n\n\n```ocaml\n(* Run 2 consumes the artifact produced by Run 1. *)\nlet store = Store.open_ () in\nmatch Store.find_artifact store ~name:\"features\" ~version:\"latest\" with\n| Some artifact ->\n Session.use_artifact session artifact;\n let path = Artifact.path artifact in\n Printf.printf \"loading from: %s\\n\" path\n| None ->\n failwith \"artifact not found\"\n```\n\nAfter this, the lineage is recorded in both directions:\n\n\n```ocaml\nArtifact.producer_run_id artifact (* run that created it *)\nArtifact.consumer_run_ids artifact (* runs that consumed it *)\n\nRun.output_artifacts run1 (* artifacts produced by run1 *)\nRun.input_artifacts run2 (* artifacts consumed by run2 *)\n```\n\n## Loading Artifacts\n\n### From a Store\n\n`Store.find_artifact` resolves by name and version (or alias):\n\n\n```ocaml\nlet store = Store.open_ () in\nmatch Store.find_artifact store ~name:\"mnist-cnn\" ~version:\"best\" with\n| Some artifact ->\n let path = Artifact.path artifact in\n Printf.printf \"path: %s (%d bytes)\\n\" path (Artifact.size_bytes artifact)\n| None ->\n Printf.printf \"not found\\n\"\n```\n\n### Listing Artifacts\n\n`Store.list_artifacts` supports filtering by name, kind, alias, and\nlineage:\n\n\n```ocaml\nlet store = Store.open_ () in\n\n(* All artifacts. *)\nlet all = Store.list_artifacts store () in\n\n(* Only checkpoints. *)\nlet checkpoints = Store.list_artifacts store ~kind:`checkpoint () in\n\n(* Only artifacts produced by a specific run. *)\nlet from_run = Store.list_artifacts store ~producer_run:\"\" () in\n\nPrintf.printf \"total: %d, checkpoints: %d, from run: %d\\n\"\n (List.length all) (List.length checkpoints) (List.length from_run)\n```\n\n### From the CLI\n\n\n```sh\n# List all artifacts.\nmunin artifacts\n\n# Filter by name.\nmunin artifacts --name mnist-cnn\n```\n\n## Complete Example: Cross-Run Lineage\n\nA data-preparation run produces a dataset. A training run consumes it\nand produces a model checkpoint.\n\n\n```ocaml\nopen Munin\n\nlet write_file path text =\n let oc = open_out path in\n Fun.protect ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc text)\n\nlet () =\n (* Run 1: produce a dataset. *)\n Session.with_run ~experiment:\"pipeline\" ~name:\"prepare-data\"\n ~tags:[ \"data\" ]\n @@ fun session ->\n write_file \"/tmp/data.csv\" \"x,y\\n1.0,2.0\\n3.0,4.0\\n\";\n ignore\n (Session.log_artifact session ~name:\"training-data\" ~kind:`dataset\n ~path:\"/tmp/data.csv\"\n ~metadata:[ (\"rows\", `Int 2) ]\n ~aliases:[ \"latest\" ] ())\n\nlet () =\n (* Run 2: consume the dataset, produce a model. *)\n let store = Store.open_ () in\n Session.with_run ~experiment:\"pipeline\" ~name:\"train\"\n ~tags:[ \"training\" ]\n @@ fun session ->\n let dataset =\n match Store.find_artifact store ~name:\"training-data\" ~version:\"latest\" with\n | Some a -> a\n | None -> failwith \"dataset not found\"\n in\n Session.use_artifact session dataset;\n Printf.printf \"training on: %s\\n\" (Artifact.path dataset);\n\n (* ... train model ... *)\n\n write_file \"/tmp/model.bin\" \"model weights\";\n ignore\n (Session.log_artifact session ~name:\"my-model\" ~kind:`model\n ~path:\"/tmp/model.bin\"\n ~aliases:[ \"latest\" ] ())\n```\n\n## Garbage Collection\n\n`Store.gc` removes blobs not referenced by any artifact version:\n\n\n```ocaml\nlet store = Store.open_ () in\nlet removed = Store.gc store in\nPrintf.printf \"removed %d unreferenced blobs\\n\" removed\n```\n"
},
{
"path": "packages/munin/doc/04-cli.md",
"content": "# CLI Reference\n\nThe `munin` command-line tool inspects and manages the local tracking store.\nEvery subcommand accepts `--root DIR` to override the default store location\n(`$RAVEN_TRACKING_DIR` or `$XDG_DATA_HOME/raven/munin`).\n\n## munin runs\n\nList tracked runs. Output is tab-separated: ID, experiment, status, parent,\nname, git commit.\n\n\n```\nmunin runs\n```\n\n```\n20260317T143201_abc mnist-sweep finished - lr-0.001 a1b2c3d\n20260317T141502_def mnist-sweep finished - lr-0.01 a1b2c3d\n20260317T140003_ghi cifar10 running - baseline e4f5a6b\n```\n\nFilter by experiment:\n\n\n```\nmunin runs --experiment mnist-sweep\n```\n\n```\n20260317T143201_abc mnist-sweep finished - lr-0.001 a1b2c3d\n20260317T141502_def mnist-sweep finished - lr-0.01 a1b2c3d\n```\n\n## munin show\n\nDisplay full details for a single run.\n\n\n```\nmunin show 20260317T143201_abc\n```\n\n```\nid: 20260317T143201_abc\nexperiment: mnist-sweep\nname: lr-0.001\nparent: -\nstatus: finished\nstarted_at: 1742224321\nended_at: 1742225180\nresumable: false\nnotes: -\ncommand: ./train.exe --lr 0.001\ncwd: /home/user/project\nhostname: workstation\npid: 42519\ngit_commit: a1b2c3d\ngit_dirty: false\nenv:\n CUDA_VISIBLE_DEVICES=0\ntags:\n sweep\n final\nparams:\n lr: 0.001\n batch_size: 64\n epochs: 10\nsummary:\n loss: 0.0312\n accuracy: 0.991\nmetric_keys:\n accuracy, loss, lr\nlatest_metrics:\n accuracy: step=9380 value=0.991\n loss: step=9380 value=0.0312\nchildren:\noutput_artifacts:\n mnist-model v3 aliases=[latest] consumers=[]\ninput_artifacts:\n mnist-data v1 producer=20260310T090000_xyz\n```\n\n## munin compare\n\nCompare two or more runs side by side. Prints a tab-separated table with\nparameters and summary values. When a metric has a declared goal (`Minimize` or\n`Maximize`), the best value is marked with `*`.\n\n\n```\nmunin compare 20260317T143201_abc 20260317T141502_def\n```\n\n```\nkey lr-0.001 lr-0.01\nbatch_size 64 64\nepochs 10 10\nlr 0.001 0.01\naccuracy 0.991* 0.984\nloss 0.0312* 0.0587\n```\n\nWorks with any number of runs:\n\n\n```\nmunin compare abc def ghi\n```\n\n## munin metrics\n\nTwo modes: listing mode (no `--key`) and history mode (with `--key`).\n\n### Listing mode\n\nShows all metric keys with their latest value, latest step, and sample count.\n\n\n```\nmunin metrics 20260317T143201_abc\n```\n\n```\nkey latest_value latest_step count\naccuracy 0.991 9380 9380\nloss 0.0312 9380 9380\nlr 0.001 9380 1\nsys/cpu_user 12.3 627 627\nsys/mem_used_gb 6.82 627 627\n```\n\n### History mode\n\nDump the full time series for a single key. Supports `--format tsv` (default),\n`csv`, and `json`.\n\n\n```\nmunin metrics 20260317T143201_abc --key loss\n```\n\n```\nstep\ttimestamp\tvalue\n1\t1742224322.123456\t2.3026\n2\t1742224322.234567\t1.8451\n3\t1742224322.345678\t1.2107\n...\n```\n\n\n```\nmunin metrics 20260317T143201_abc --key loss --format csv\n```\n\n```\nstep,timestamp,value\n1,1742224322.123456,2.3026\n2,1742224322.234567,1.8451\n...\n```\n\n\n```\nmunin metrics 20260317T143201_abc --key loss --format json\n```\n\n```json\n[{\"step\":1,\"timestamp\":1742224322.123456,\"value\":2.3026},{\"step\":2,\"timestamp\":1742224322.234567,\"value\":1.8451}]\n```\n\n## munin watch\n\nLaunch the terminal dashboard. See [Terminal Dashboard](05-dashboard.md) for\nfull documentation.\n\nAuto-detect the latest run:\n\n\n```\nmunin watch\n```\n\nOpen a specific run:\n\n\n```\nmunin watch 20260317T143201_abc\n```\n\nFilter by experiment (picks the latest run in that experiment):\n\n\n```\nmunin watch --experiment mnist-sweep\n```\n\n## munin artifacts\n\nList stored artifacts. Output is tab-separated: name, version, kind, payload\ntype, size in bytes, producer run, consumer runs.\n\n\n```\nmunin artifacts\n```\n\n```\nmnist-data v1 dataset dir 48000000 20260310T090000_xyz 20260317T143201_abc,20260317T141502_def\nmnist-model v1 model file 4521984 20260317T141502_def -\nmnist-model v2 model file 4521984 20260317T143201_abc -\nmnist-model v3 model file 4521984 20260317T143201_abc -\n```\n\nFilter by name:\n\n\n```\nmunin artifacts --name mnist-model\n```\n\n```\nmnist-model v1 model file 4521984 20260317T141502_def -\nmnist-model v2 model file 4521984 20260317T143201_abc -\nmnist-model v3 model file 4521984 20260317T143201_abc -\n```\n\n## munin delete\n\nDelete a run and its event log from the store. Does not remove shared blobs\n(use `munin gc` for that). Removes the experiment directory if no runs remain.\n\n\n```\nmunin delete 20260317T141502_def\n```\n\n```\nDelete run 20260317T141502_def (mnist-sweep / lr-0.01)? [y/N] y\nDeleted.\n```\n\nSkip the confirmation prompt with `--yes`:\n\n\n```\nmunin delete 20260317T141502_def --yes\n```\n\n## munin gc\n\nGarbage-collect unreferenced blobs from the blob store. Blobs that are no\nlonger referenced by any artifact are removed.\n\n\n```\nmunin gc\n```\n\n```\nRemoved 3 unreferenced blob(s).\n```\n"
},
{
"path": "packages/munin/doc/05-dashboard.md",
"content": "# Terminal Dashboard\n\nMunin includes a terminal-based dashboard for monitoring runs in real time.\nIt renders braille-resolution charts, status indicators, and system resource\nbars directly in the terminal.\n\n## Launching\n\nThe dashboard is started with `munin watch`. With no arguments it auto-detects\nthe most recently started run:\n\n\n```\nmunin watch\n```\n\nTo open a specific run, pass its ID:\n\n\n```\nmunin watch 20260317T143201_abc\n```\n\nTo pick the latest run in a given experiment:\n\n\n```\nmunin watch --experiment mnist-sweep\n```\n\n## Layout\n\nThe dashboard has three sections stacked vertically:\n\n### Header\n\nA single-line bar showing:\n\n- **Experiment and run name** (with run ID in parentheses)\n- **Tags** as inline badges\n- **Epoch** counter (e.g. `Epoch 3/10`) when an `epoch` metric is logged and\n `epochs` is set in the params\n- **Step** counter (the highest step across all metrics)\n- **Elapsed time** in `HH:MM:SS`\n- **Status badge** on the right: a colored dot and label\n\n### Metrics panel\n\nA grid of braille-resolution line charts, one per user metric (system metrics\nprefixed with `sys/` are excluded). Each chart shows the metric name, latest\nvalue, and best value when a goal is defined.\n\nCharts are arranged in a responsive grid: 2 columns when the terminal is wide\nenough (at least 50 characters), 1 column otherwise. Rows are sized at 14\ncharacters tall. When there are more metrics than fit on screen, they are split\ninto batches and navigated with `<` / `>`.\n\nThe currently selected chart has a white border; unselected charts have a dim\nborder. Pressing Enter on the selected chart opens the detail view.\n\n### System panel\n\nA side panel (right 34% of the screen) showing four resource bars when system\nmetrics are available:\n\n- **CPU** -- combined user (green) + system (cyan) percentage with sparkline\n- **Mem** -- system memory percentage and absolute GB with sparkline\n- **Proc** -- process CPU percentage with sparkline\n- **RSS** -- process resident set size in MB with sparkline\n\nThe bars change color based on utilization: green below 50%, yellow 50-80%,\nred above 80%.\n\nToggle the system panel on/off with `[` or `]`.\n\n### Footer\n\nA hint bar showing available keyboard shortcuts for the current mode.\n\n## Keyboard shortcuts\n\n### Dashboard mode\n\n| Key | Action |\n|--------------|--------------------------------------------|\n| Arrow keys | Navigate the metric chart grid |\n| Enter/Space | Open the selected metric in detail view |\n| `<` / `>` | Previous / next batch of metrics |\n| `[` / `]` | Toggle the system panel |\n| `q` / Escape | Quit the dashboard |\n\n### Detail view\n\n| Key | Action |\n|--------------|--------------------------------------------|\n| `S` | Cycle EMA smoothing: Off, Light (1), Medium (2), Heavy (3) |\n| `q` / Escape | Return to dashboard |\n\n## Status detection\n\nThe dashboard determines run status from the event log:\n\n- **Live** (green) -- new events are arriving.\n- **Stopped** (gray) -- no events received for 5 seconds. The run process may\n have crashed or been suspended.\n- **Done** (blue) -- a `Finished` event with status `finished` was received.\n- **Failed** (red) -- a `Finished` event with status `failed` was received.\n- **Killed** (yellow) -- a `Finished` event with status `killed` was received.\n\nThe dashboard polls the event log on every tick and transitions between states\nautomatically.\n\n## Detail view\n\nPressing Enter on a metric chart opens a full-screen detail view. The chart\nfills 80% of the screen with full axis labels and gridlines.\n\n**EMA smoothing** can be toggled by pressing `S`, cycling through four levels:\n\n| Level | Alpha | Effect |\n|--------|-------|-------------------------------|\n| Off | -- | Raw values |\n| Light | 0.5 | Mild smoothing |\n| Medium | 0.3 | Moderate smoothing |\n| Heavy | 0.15 | Aggressive smoothing |\n\nWhen smoothing is active, the chart title shows \"(EMA)\" and the footer displays\nthe current level number.\n\nThe best value (determined by the metric's declared goal, or heuristically for\nkeys containing \"loss\" or \"error\") is displayed below the chart.\n"
},
{
"path": "packages/munin/doc/06-system-monitoring.md",
"content": "# System Monitoring\n\nThe `munin.sys` library provides background system monitoring that logs CPU and\nmemory metrics alongside your training metrics.\n\n## Setup\n\nAdd `munin.sys` to your dune `libraries`:\n\n\n```lisp\n(executable\n (name main)\n (libraries munin munin.sys))\n```\n\n## Usage\n\nStart a monitor after creating a session, and stop it before finishing:\n\n\n```ocaml\nlet () =\n let session =\n Munin.Session.start ~experiment:\"train\" ~name:\"resnet-50\" ()\n in\n let monitor = Munin_sys.start session () in\n\n (* ... training loop ... *)\n\n Munin_sys.stop monitor;\n Munin.Session.finish session ()\n```\n\nThe monitor spawns a background thread that samples system and process\nstatistics at a fixed interval and logs them as scalar metrics.\n\n### Configuring the interval\n\nThe default sampling interval is 15 seconds. Override it with `~interval`:\n\n\n```ocaml\nlet monitor = Munin_sys.start session ~interval:5.0 ()\n```\n\nThe first sample is taken after one interval elapses.\n\n## Logged metrics\n\nAll metrics use the `sys/` prefix and are defined with `~summary:`Last`, so the\nfinal sampled value appears in run summaries.\n\n### System-wide\n\n| Metric key | Description | Range |\n|---------------------|-----------------------------------|----------|\n| `sys/cpu_user` | User CPU percentage | 0--100 |\n| `sys/cpu_system` | System (kernel) CPU percentage | 0--100 |\n| `sys/mem_used_pct` | Memory usage percentage | 0--100 |\n| `sys/mem_used_gb` | Memory used in GB | 0+ |\n\n### Per-process\n\n| Metric key | Description | Range |\n|---------------------|-----------------------------------|----------|\n| `sys/proc_cpu_pct` | Process CPU percentage | 0+ |\n| `sys/proc_mem_mb` | Process resident set size in MB | 0+ |\n\n## Platform support\n\n`munin.sys` works on Linux and macOS. Platform-specific behavior:\n\n- **Linux**: CPU counters are read from `/proc/stat`; memory from\n `/proc/meminfo`; process stats from `/proc/self/stat` and `Unix.times`.\n- **macOS**: CPU and memory use Mach host statistics; process memory uses\n `task_info`. Only user/nice/system/idle CPU fields are populated.\n\n## TUI system panel\n\nThe `munin watch` dashboard displays a system panel on the right side with\nCPU, memory, process CPU, and RSS bars. This panel reads the `sys/` metrics\nfrom the run's event log -- it does not perform its own sampling. If your run\ndoes not use `munin.sys`, the system panel shows zeroes.\n\nToggle the panel with `[` or `]` in the dashboard.\n\n## Sysstat module\n\nThe `Munin_sys` module re-exports the `Sysstat` interface, giving direct access\nto stateless, poll-based sampling functions. These are useful for custom\nmonitoring outside the background thread.\n\n### Cpu\n\n\n```ocaml\nlet prev = Munin_sys.Cpu.sample () in\n(* ... wait ... *)\nlet next = Munin_sys.Cpu.sample () in\nlet stats = Munin_sys.Cpu.compute ~prev ~next in\nPrintf.printf \"User: %.1f%% System: %.1f%%\\n\" stats.user stats.system\n```\n\n`Cpu.sample_per_core` returns an array of per-core counters.\n\n### Mem\n\n\n```ocaml\nlet mem = Munin_sys.Mem.sample () in\nlet used_gb = Int64.to_float mem.used /. 1_073_741_824. in\nPrintf.printf \"Memory: %.1f GB used / %.1f GB total\\n\"\n used_gb (Int64.to_float mem.total /. 1_073_741_824.)\n```\n\n### Net\n\n\n```ocaml\nlet prev = Munin_sys.Net.sample () in\n(* ... wait ... *)\nlet next = Munin_sys.Net.sample () in\nlet stats = Munin_sys.Net.compute ~prev ~next ~dt:1.0 in\nPrintf.printf \"Rx: %.0f B/s Tx: %.0f B/s\\n\"\n stats.rx_bytes_per_sec stats.tx_bytes_per_sec\n```\n\n### Disk_io\n\n\n```ocaml\nlet prev = Munin_sys.Disk_io.sample () in\n(* ... wait ... *)\nlet next = Munin_sys.Disk_io.sample () in\nlet stats = Munin_sys.Disk_io.compute ~prev ~next ~dt:1.0 in\nPrintf.printf \"Read: %.0f B/s Write: %.0f B/s Util: %.1f%%\\n\"\n stats.read_bytes_per_sec stats.write_bytes_per_sec\n stats.utilization_percent\n```\n\n### Fs\n\n\n```ocaml\nlet fs = Munin_sys.Fs.sample () in\nlet used_pct =\n Int64.to_float fs.used_bytes /. Int64.to_float fs.total_bytes *. 100.\nin\nPrintf.printf \"Disk: %.1f%% used\\n\" used_pct;\nList.iter (fun (p : Munin_sys.Fs.partition) ->\n Printf.printf \" %s: %Ld / %Ld bytes\\n\"\n p.mount_point p.used_bytes p.total_bytes\n) fs.partitions\n```\n\n### Proc\n\nCurrent process stats:\n\n\n```ocaml\nlet prev = Munin_sys.Proc.Self.sample () in\n(* ... wait ... *)\nlet next = Munin_sys.Proc.Self.sample () in\nlet stats = Munin_sys.Proc.Self.compute ~prev ~next ~dt:1.0 ~num_cores:None in\nPrintf.printf \"CPU: %.1f%% RSS: %Ld bytes\\n\" stats.cpu_percent stats.rss_bytes\n```\n\nProcess table (all visible processes):\n\n\n```ocaml\nlet prev = Munin_sys.Proc.Table.sample () in\n(* ... wait ... *)\nlet next = Munin_sys.Proc.Table.sample () in\nlet stats = Munin_sys.Proc.Table.compute ~prev ~next ~dt:1.0 in\nList.iter (fun (p : Munin_sys.Proc.Table.stats) ->\n Printf.printf \"%d %-15s CPU: %.1f%% Mem: %.1f%%\\n\"\n p.pid p.name p.cpu_percent p.mem_percent\n) (List.sort (fun a b -> compare b.cpu_percent a.cpu_percent) stats)\n```\n\n### System info\n\n\n```ocaml\nlet (l1, l5, l15) = Munin_sys.loadavg () in\nPrintf.printf \"Load: %.2f %.2f %.2f\\n\" l1 l5 l15;\nPrintf.printf \"Uptime: %Ld seconds\\n\" (Munin_sys.uptime ())\n```\n"
},
{
"path": "packages/munin/doc/dune",
"content": "(mdx\n (files *.md)\n (package munin)\n (deps))\n"
},
{
"path": "packages/munin/doc/index.md",
"content": "# Munin\n\nMunin is a local-first experiment tracker for OCaml. It records\nhyperparameters, metrics, media, and versioned artifacts to disk with\nno external services. A CLI and live TUI let you inspect and compare\nruns from the terminal.\n\n## Features\n\n- **Scalar tracking**: `log_metric`, `log_metrics`, auto-computed summaries\n- **Metric definitions**: summary modes (min/max/mean/last), goals (minimize/maximize), custom x-axes\n- **Media logging**: images, files, audio, and structured tables\n- **Versioned artifacts**: content-addressed deduplication, aliases, cross-run lineage\n- **Provenance**: git commit, command line, environment variables, captured automatically\n- **System monitoring**: background CPU and memory sampling via `Munin_sys`\n- **CLI**: `munin runs`, `munin show`, `munin compare`, `munin metrics`, `munin artifacts`\n- **Live TUI**: `munin watch` with real-time metric charts and system stats\n\n## Quick Start\n\n\n```ocaml\nlet () =\n Munin.Session.with_run ~experiment:\"demo\"\n ~params:[ (\"lr\", `Float 0.001); (\"epochs\", `Int 10) ]\n @@ fun session ->\n for step = 1 to 100 do\n let loss = 1.0 /. Float.of_int step in\n Munin.Session.log_metric session ~step \"loss\" loss\n done;\n Munin.Session.set_summary session [ (\"final_loss\", `Float 0.01) ]\n```\n\nInspect the run from the terminal:\n\n\n```sh\nmunin runs\nmunin show \nmunin metrics --key loss\n```\n\n## Next Steps\n\n- [Getting Started](01-getting-started/) -- installation, key concepts, first example\n- [Tracking Metrics](02-tracking/) -- scalars, metric definitions, media, Kaun integration\n- [Artifacts](03-artifacts/) -- versioned files, aliases, lineage, deduplication\n"
},
{
"path": "packages/munin/examples/01-basic/README.md",
"content": "# 01-basic\n\nCreates a small local store, logs a run with scalar metrics, writes a summary,\nstores a file artifact, and prints the resulting run id.\n"
},
{
"path": "packages/munin/examples/01-basic/dune",
"content": "(executable\n (name main)\n (libraries munin))\n"
},
{
"path": "packages/munin/examples/01-basic/main.ml",
"content": "let () =\n let root = \"_munin\" in\n let artifact_path = Filename.concat root \"artifact.txt\" in\n let write path text =\n let oc = open_out path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc text)\n in\n let session =\n Session.start ~root ~experiment:\"demo\" ~name:\"baseline\"\n ~params:[ (\"lr\", `Float 0.001) ]\n ()\n in\n write artifact_path \"hello from munin\\n\";\n Session.log_metric session ~step:1 \"loss\" 1.25;\n Session.log_metric session ~step:2 \"loss\" 0.94;\n Session.set_summary session [ (\"best_loss\", `Float 0.94) ];\n ignore\n (Session.log_artifact session ~name:\"notes\" ~kind:`file ~path:artifact_path\n ());\n Session.finish session ();\n Printf.printf \"run: %s\\n\" (Run.id (Session.run session))\n"
},
{
"path": "packages/munin/examples/02-metrics/dune",
"content": "(executable\n (name main)\n (libraries munin))\n"
},
{
"path": "packages/munin/examples/02-metrics/main.ml",
"content": "(** Metric definitions and rich scalar logging.\n\n Demonstrates define_metric with summaries, goals, and step_metric for custom\n x-axes. Simulates an iterative solver converging over epochs. *)\n\nopen Munin\n\nlet () =\n let root = \"_munin\" in\n let session =\n Session.start ~root ~experiment:\"solver\" ~name:\"conjugate-gradient\"\n ~params:[ (\"tolerance\", `Float 1e-6); (\"max_iter\", `Int 500) ]\n ()\n in\n (* Declare how metrics should be summarised and compared. *)\n Session.define_metric session \"residual\" ~summary:`Min ~goal:`Minimize ();\n Session.define_metric session \"convergence_rate\" ~summary:`Mean\n ~step_metric:\"epoch\" ();\n\n (* Simulate an iterative solver: residual shrinks, rate stabilises. *)\n let residual = ref 1.0 in\n for epoch = 1 to 20 do\n let rate = 0.7 +. Random.float 0.1 in\n residual := !residual *. rate;\n let step = epoch * 25 in\n Session.log_metrics session ~step\n [\n (\"residual\", !residual);\n (\"convergence_rate\", rate);\n (\"epoch\", Float.of_int epoch);\n ]\n done;\n\n Session.set_summary session [ (\"final_residual\", `Float !residual) ];\n Session.finish session ();\n\n (* Read back and print. *)\n let run = Session.run session in\n Printf.printf \"run: %s\\n\" (Run.id run);\n Printf.printf \"metric keys: %s\\n\" (String.concat \", \" (Run.metric_keys run));\n\n let defs = Run.metric_defs run in\n List.iter\n (fun (key, (def : Run.metric_def)) ->\n let goal =\n match def.goal with\n | Some `Minimize -> \"minimize\"\n | Some `Maximize -> \"maximize\"\n | None -> \"none\"\n in\n Printf.printf \" %s: summary=%s goal=%s\\n\" key\n (match def.summary with\n | `Min -> \"min\"\n | `Max -> \"max\"\n | `Mean -> \"mean\"\n | `Last -> \"last\"\n | `None -> \"none\")\n goal)\n defs;\n\n let history = Run.metric_history run \"residual\" in\n Printf.printf \"residual: %d samples, final=%.2e\\n\" (List.length history)\n (List.nth history (List.length history - 1)).value\n"
},
{
"path": "packages/munin/examples/03-artifacts/dune",
"content": "(executable\n (name main)\n (libraries munin))\n"
},
{
"path": "packages/munin/examples/03-artifacts/main.ml",
"content": "(** Artifact versioning and lineage across runs.\n\n Run 1 produces a dataset artifact. Run 2 consumes it and produces a result.\n Demonstrates versioning, aliases, and cross-run lineage. *)\n\nopen Munin\n\nlet write_file path text =\n let oc = open_out path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc text)\n\nlet () =\n let root = \"_munin\" in\n let store = Store.open_ ~root () in\n\n (* Run 1: produce a dataset. *)\n let session1 =\n Session.start ~root ~experiment:\"pipeline\" ~name:\"prepare-data\"\n ~tags:[ \"data\" ] ()\n in\n let data_path = Filename.concat root \"measurements.csv\" in\n write_file data_path \"wavelength,flux\\n450.0,1.23\\n550.0,2.45\\n650.0,1.87\\n\";\n let dataset =\n Session.log_artifact session1 ~name:\"measurements\" ~kind:`dataset\n ~path:data_path\n ~metadata:[ (\"rows\", `Int 3) ]\n ~aliases:[ \"latest\" ] ()\n in\n Session.finish session1 ();\n Printf.printf \"produced: %s v%s (aliases: %s)\\n\" (Artifact.name dataset)\n (Artifact.version dataset)\n (String.concat \", \" (Artifact.aliases dataset));\n\n (* Run 2: consume the dataset, produce a result. *)\n let session2 =\n Session.start ~root ~experiment:\"pipeline\" ~name:\"analyse\"\n ~tags:[ \"analysis\" ] ()\n in\n Session.use_artifact session2 dataset;\n let result_path = Filename.concat root \"result.txt\" in\n write_file result_path \"peak_wavelength=550.0\\npeak_flux=2.45\\n\";\n let result =\n Session.log_artifact session2 ~name:\"analysis-result\" ~kind:`file\n ~path:result_path ~aliases:[ \"latest\"; \"best\" ] ()\n in\n Session.finish session2 ();\n Printf.printf \"produced: %s v%s\\n\" (Artifact.name result)\n (Artifact.version result);\n\n (* Query artifacts from the store. *)\n (match Store.find_artifact store ~name:\"measurements\" ~version:\"latest\" with\n | Some a ->\n Printf.printf \"\\nresolved 'measurements:latest' -> v%s (%d bytes)\\n\"\n (Artifact.version a) (Artifact.size_bytes a);\n Printf.printf \" producer: %s\\n\"\n (Option.value ~default:\"unknown\" (Artifact.producer_run_id a));\n Printf.printf \" consumers: %s\\n\"\n (String.concat \", \" (Artifact.consumer_run_ids a))\n | None -> Printf.printf \"artifact not found\\n\");\n\n let all = Store.list_artifacts store () in\n Printf.printf \"\\nall artifacts: %d\\n\" (List.length all);\n List.iter\n (fun a -> Printf.printf \" %s v%s\\n\" (Artifact.name a) (Artifact.version a))\n all\n"
},
{
"path": "packages/munin/examples/04-media/dune",
"content": "(executable\n (name main)\n (libraries munin unix))\n"
},
{
"path": "packages/munin/examples/04-media/main.ml",
"content": "(** Non-scalar data: media files and tables.\n\n Demonstrates log_media for files and log_table for structured data. Creates\n synthetic data to keep the example self-contained. *)\n\nopen Munin\n\nlet write_file path text =\n let oc = open_out path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc text)\n\n(* Write a tiny PPM image (3x3 red gradient). *)\nlet write_ppm path =\n let oc = open_out_bin path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () ->\n Printf.fprintf oc \"P6\\n3 3\\n255\\n\";\n for row = 0 to 2 do\n for _col = 0 to 2 do\n let v = 85 * row in\n output_char oc (Char.chr v);\n output_char oc '\\000';\n output_char oc '\\000'\n done\n done)\n\nlet () =\n let root = \"_munin\" in\n let session = Session.start ~root ~experiment:\"media-demo\" ~name:\"run-1\" () in\n let tmp = Filename.concat root \"_tmp\" in\n (try Unix.mkdir tmp 0o755 with Unix.Unix_error (Unix.EEXIST, _, _) -> ());\n\n (* Log an image at two different steps. *)\n let img_path = Filename.concat tmp \"sample.ppm\" in\n write_ppm img_path;\n Session.log_media session ~step:1 ~key:\"viz/sample\" ~kind:`Image\n ~path:img_path;\n write_ppm img_path;\n Session.log_media session ~step:2 ~key:\"viz/sample\" ~kind:`Image\n ~path:img_path;\n\n (* Log a text file. *)\n let notes_path = Filename.concat tmp \"notes.txt\" in\n write_file notes_path \"Observation: signal peaks at 550nm.\\n\";\n Session.log_media session ~step:1 ~key:\"notes\" ~kind:`File ~path:notes_path;\n\n (* Log a structured table — e.g. per-class metrics or a confusion matrix. *)\n Session.log_table session ~step:1 ~key:\"results/per_band\"\n ~columns:[ \"band\"; \"snr\"; \"coverage\" ]\n ~rows:\n [\n [ `String \"blue\"; `Float 12.3; `Float 0.95 ];\n [ `String \"green\"; `Float 18.7; `Float 0.98 ];\n [ `String \"red\"; `Float 15.1; `Float 0.92 ];\n ];\n\n Session.finish session ();\n\n (* Read back media entries. *)\n let run = Session.run session in\n Printf.printf \"run: %s\\n\" (Run.id run);\n Printf.printf \"media keys: %s\\n\" (String.concat \", \" (Run.media_keys run));\n let entries = Run.media_history run \"viz/sample\" in\n Printf.printf \"viz/sample: %d entries\\n\" (List.length entries);\n List.iter\n (fun (e : Run.media_entry) ->\n Printf.printf \" step=%d kind=%s path=%s\\n\" e.step\n (match e.kind with\n | `Image -> \"image\"\n | `Audio -> \"audio\"\n | `Table -> \"table\"\n | `File -> \"file\")\n (Filename.basename e.path))\n entries\n"
},
{
"path": "packages/munin/examples/05-parameter-sweep/dune",
"content": "(executable\n (name main)\n (libraries munin))\n"
},
{
"path": "packages/munin/examples/05-parameter-sweep/main.ml",
"content": "(** Parameter sweep with grouped runs.\n\n Runs the same computation with different configurations, grouped under a\n single sweep. Compares results at the end via Store queries. *)\n\nopen Munin\n\nlet () =\n let root = \"_munin\" in\n let store = Store.open_ ~root () in\n let group = \"sweep-1\" in\n\n (* Sweep over three configurations. *)\n let configs =\n [\n (\"aggressive\", 0.1, 50);\n (\"moderate\", 0.01, 100);\n (\"conservative\", 0.001, 200);\n ]\n in\n List.iter\n (fun (name, step_size, max_iter) ->\n let session =\n Session.start ~root ~experiment:\"optimisation\" ~name ~group\n ~params:\n [ (\"step_size\", `Float step_size); (\"max_iter\", `Int max_iter) ]\n ()\n in\n Session.define_metric session \"error\" ~summary:`Min ~goal:`Minimize ();\n\n (* Simulate convergence: smaller steps converge slower but lower. *)\n let error = ref 10.0 in\n for i = 1 to max_iter do\n error := (!error *. (1.0 -. step_size)) +. Random.float 0.01;\n if i mod 10 = 0 then Session.log_metric session ~step:i \"error\" !error\n done;\n Session.finish session ())\n configs;\n\n (* Compare: list all runs in the group and print a results table. *)\n let runs = Store.list_runs store ~experiment:\"optimisation\" ~group () in\n Printf.printf \"%-15s %-10s %-10s %-12s\\n\" \"name\" \"step_size\" \"max_iter\"\n \"final_error\";\n Printf.printf \"%s\\n\" (String.make 52 '-');\n List.iter\n (fun run ->\n let name = Option.value ~default:\"?\" (Run.name run) in\n let step_size =\n match Run.find_param run \"step_size\" with\n | Some (`Float f) -> Printf.sprintf \"%g\" f\n | _ -> \"?\"\n in\n let max_iter =\n match Run.find_param run \"max_iter\" with\n | Some v -> Format.asprintf \"%a\" Value.pp v\n | None -> \"?\"\n in\n let latest = Run.latest_metrics run in\n let error =\n match List.assoc_opt \"error\" latest with\n | Some m -> Printf.sprintf \"%.6f\" m.value\n | None -> \"?\"\n in\n Printf.printf \"%-15s %-10s %-10s %-12s\\n\" name step_size max_iter error)\n runs\n"
},
{
"path": "packages/munin/examples/06-inspect/dune",
"content": "(executable\n (name main)\n (libraries munin))\n"
},
{
"path": "packages/munin/examples/06-inspect/main.ml",
"content": "(** Querying and inspecting past runs.\n\n The \"notebook\" use case: open a store, browse experiments, filter runs,\n examine provenance, and extract metric histories. Assumes earlier examples\n have been run to populate the store. *)\n\nopen Munin\n\nlet () =\n let root = \"_munin\" in\n let store = Store.open_ ~root () in\n\n (* List all experiments. *)\n let experiments = Store.list_experiments store in\n Printf.printf \"experiments: %s\\n\\n\" (String.concat \", \" experiments);\n\n (* List runs, optionally filtering. *)\n let all_runs = Store.list_runs store () in\n Printf.printf \"total runs: %d\\n\" (List.length all_runs);\n let finished = Store.list_runs store ~status:`finished () in\n Printf.printf \"finished runs: %d\\n\\n\" (List.length finished);\n\n (* Find the latest run and inspect it. *)\n (match Store.latest_run store () with\n | None -> Printf.printf \"no runs found\\n\"\n | Some run ->\n Printf.printf \"latest run: %s\\n\" (Run.id run);\n Printf.printf \" experiment: %s\\n\" (Run.experiment_name run);\n Printf.printf \" name: %s\\n\"\n (Option.value ~default:\"(none)\" (Run.name run));\n Printf.printf \" status: %s\\n\"\n (match Run.status run with\n | `running -> \"running\"\n | `finished -> \"finished\"\n | `failed -> \"failed\"\n | `killed -> \"killed\");\n Printf.printf \" tags: [%s]\\n\" (String.concat \", \" (Run.tags run));\n\n (* Provenance. *)\n let prov = Run.provenance run in\n Printf.printf \" hostname: %s\\n\" (Option.value ~default:\"?\" prov.hostname);\n Printf.printf \" git: %s%s\\n\"\n (Option.value ~default:\"?\" prov.git_commit)\n (match prov.git_dirty with Some true -> \" (dirty)\" | _ -> \"\");\n\n (* Params. *)\n let params = Run.params run in\n if params <> [] then (\n Printf.printf \" params:\\n\";\n List.iter\n (fun (k, v) ->\n Printf.printf \" %s = %s\\n\" k (Format.asprintf \"%a\" Value.pp v))\n params);\n\n (* Metrics. *)\n let keys = Run.metric_keys run in\n Printf.printf \" metrics: %s\\n\" (String.concat \", \" keys);\n List.iter\n (fun key ->\n let history = Run.metric_history run key in\n let n = List.length history in\n if n > 0 then\n let last = (List.nth history (n - 1)).value in\n Printf.printf \" %s: %d samples, last=%.4g\\n\" key n last)\n keys);\n\n (* List artifacts. *)\n let artifacts = Store.list_artifacts store () in\n Printf.printf \"\\nartifacts: %d\\n\" (List.length artifacts);\n List.iter\n (fun a ->\n Printf.printf \" %s v%s (%d bytes)\\n\" (Artifact.name a)\n (Artifact.version a) (Artifact.size_bytes a))\n artifacts\n"
},
{
"path": "packages/munin/examples/07-system-monitor/dune",
"content": "(executable\n (name main)\n (libraries munin munin.sys))\n"
},
{
"path": "packages/munin/examples/07-system-monitor/main.ml",
"content": "(** Automatic system metrics during a computation.\n\n Starts a system monitor that logs CPU and memory usage in the background\n while a CPU-intensive computation runs. *)\n\nopen Munin\n\n(* A simple CPU-bound computation: count primes up to n. *)\nlet count_primes n =\n let count = ref 0 in\n for i = 2 to n do\n let is_prime = ref true in\n let j = ref 2 in\n while !j * !j <= i && !is_prime do\n if i mod !j = 0 then is_prime := false;\n incr j\n done;\n if !is_prime then incr count\n done;\n !count\n\nlet () =\n let root = \"_munin\" in\n let session =\n Session.start ~root ~experiment:\"compute\" ~name:\"prime-sieve\"\n ~params:[ (\"limit\", `Int 5_000_000) ]\n ()\n in\n\n (* Start system monitoring with a short interval for this demo. *)\n let monitor = Munin_sys.start session ~interval:0.5 () in\n\n (* Run the computation, logging progress. *)\n let steps = 10 in\n let per_step = 500_000 in\n for i = 1 to steps do\n let limit = i * per_step in\n let n = count_primes limit in\n Session.log_metrics session ~step:i\n [ (\"primes_found\", Float.of_int n); (\"limit\", Float.of_int limit) ]\n done;\n\n Munin_sys.stop monitor;\n Session.finish session ();\n\n (* Check what system metrics were recorded. *)\n let run = Session.run session in\n let keys = Run.metric_keys run in\n let sys_keys =\n List.filter (fun k -> String.length k > 4 && String.sub k 0 4 = \"sys/\") keys\n in\n Printf.printf \"run: %s\\n\" (Run.id run);\n Printf.printf \"system metrics: %s\\n\" (String.concat \", \" sys_keys);\n List.iter\n (fun key ->\n let history = Run.metric_history run key in\n let n = List.length history in\n if n > 0 then\n let last = (List.nth history (n - 1)).value in\n Printf.printf \" %s: %d samples, last=%.2f\\n\" key n last)\n sys_keys\n"
},
{
"path": "packages/munin/examples/README.md",
"content": "# Munin Examples\n\n| Example | What you'll learn |\n|---------|-------------------|\n| [01-basic](01-basic/) | Start a session, log metrics, store an artifact, finish |\n| [02-metrics](02-metrics/) | Define metrics with summaries, goals, and custom x-axes |\n| [03-artifacts](03-artifacts/) | Version artifacts, attach aliases, track cross-run lineage |\n| [04-media](04-media/) | Log images, files, and structured tables |\n| [05-parameter-sweep](05-parameter-sweep/) | Group runs under a sweep, compare results |\n| [06-inspect](06-inspect/) | Query the store, browse experiments, examine provenance |\n| [07-system-monitor](07-system-monitor/) | Record CPU and memory usage automatically |\n| [x-kaun-mnist](x-kaun-mnist/) | End-to-end MNIST training with kaun integration |\n\nRun any example with:\n\n```sh\ndune exec packages/munin/examples/01-basic/main.exe\n```\n\nExamples write to a local `_munin/` directory.\n"
},
{
"path": "packages/munin/examples/x-kaun-mnist/dune",
"content": "(executable\n (name main)\n (libraries nx rune vega kaun kaun.datasets munin munin.sys))\n"
},
{
"path": "packages/munin/examples/x-kaun-mnist/main.ml",
"content": "(** End-to-end MNIST training with experiment tracking.\n\n Trains a CNN on MNIST using kaun, logging metrics, hyperparameters, and a\n model checkpoint via munin. Shows how munin integrates with a real training\n loop without adding a dependency from kaun to munin. *)\n\nopen Kaun\n\nlet batch_size = 64\nlet epochs = 3\nlet lr = 0.001\n\nlet model =\n Layer.sequential\n [\n Layer.conv2d ~in_channels:1 ~out_channels:16 ();\n Layer.relu ();\n Layer.max_pool2d ~kernel_size:(2, 2) ();\n Layer.conv2d ~in_channels:16 ~out_channels:32 ();\n Layer.relu ();\n Layer.max_pool2d ~kernel_size:(2, 2) ();\n Layer.flatten ();\n Layer.linear ~in_features:(32 * 7 * 7) ~out_features:128 ();\n Layer.relu ();\n Layer.linear ~in_features:128 ~out_features:10 ();\n ]\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let dtype = Nx.float32 in\n\n (* Start a tracked run. *)\n let session =\n Munin.Session.start ~experiment:\"mnist\" ~name:\"cnn-adam\"\n ~tags:[ \"baseline\" ]\n ~params:\n [\n (\"lr\", `Float lr);\n (\"batch_size\", `Int batch_size);\n (\"epochs\", `Int epochs);\n (\"optimizer\", `String \"adam\");\n ]\n ()\n in\n Munin.Session.define_metric session \"train/loss\" ~summary:`Min ~goal:`Minimize\n ();\n Munin.Session.define_metric session \"val/accuracy\" ~summary:`Max\n ~goal:`Maximize ();\n let sysmon = Munin_sys.start session () in\n\n Printf.printf \"run: %s\\n%!\" (Munin.Run.id (Munin.Session.run session));\n\n (* Load data. *)\n Printf.printf \"Loading MNIST...\\n%!\";\n let (x_train, y_train), (x_test, y_test) = Kaun_datasets.mnist () in\n let n_train = (Nx.shape x_train).(0) in\n Printf.printf \" train: %d test: %d\\n%!\" n_train (Nx.shape x_test).(0);\n\n let test_batches = Data.prepare ~batch_size (x_test, y_test) in\n let trainer =\n Train.make ~model ~optimizer:(Vega.adam (Vega.Schedule.constant lr))\n in\n let st = ref (Train.init trainer ~dtype) in\n let global_step = ref 0 in\n let last_acc = ref 0. in\n\n for epoch = 1 to epochs do\n let train_data =\n Data.prepare ~shuffle:true ~batch_size (x_train, y_train)\n |> Data.map (fun (x, y) ->\n (x, fun logits -> Loss.cross_entropy_sparse logits y))\n in\n let num_batches = n_train / batch_size in\n let tracker = Metric.tracker () in\n\n st :=\n Train.fit trainer !st\n ~report:(fun ~step ~loss _st ->\n let s = !global_step + step in\n Metric.observe tracker \"loss\" loss;\n Munin.Session.log_metrics session ~step:s\n [ (\"train/loss\", loss); (\"epoch\", Float.of_int epoch) ];\n Printf.printf \"\\r batch %d/%d loss: %.4f%!\" step num_batches loss)\n train_data;\n global_step := !global_step + num_batches;\n Printf.printf \"\\n%!\";\n\n (* Evaluate. *)\n Data.reset test_batches;\n let test_acc =\n Metric.eval\n (fun (x, y) ->\n let logits = Train.predict trainer !st x in\n Metric.accuracy logits y)\n test_batches\n in\n last_acc := test_acc;\n\n Munin.Session.log_metrics session ~step:!global_step\n [\n (\"train/loss_avg\", Metric.mean tracker \"loss\");\n (\"val/accuracy\", test_acc);\n ];\n\n Printf.printf \"epoch %d loss: %.4f val_acc: %.2f%%\\n%!\" epoch\n (Metric.mean tracker \"loss\")\n (test_acc *. 100.)\n done;\n\n (* Save model checkpoint as a versioned artifact. *)\n let checkpoint_path =\n Filename.concat\n (Munin.Run.dir (Munin.Session.run session))\n \"model.safetensors\"\n in\n Checkpoint.save checkpoint_path (Layer.params (Train.vars !st));\n ignore\n (Munin.Session.log_artifact session ~name:\"mnist-cnn\" ~kind:`checkpoint\n ~path:checkpoint_path\n ~metadata:[ (\"format\", `String \"safetensors\") ]\n ~aliases:[ \"latest\" ] ());\n\n Munin_sys.stop sysmon;\n Munin.Session.set_notes session\n (Some (Printf.sprintf \"Final val accuracy: %.2f%%\" (!last_acc *. 100.)));\n Munin.Session.finish session ();\n Printf.printf \"\\nDone. Run: %s\\n\" (Munin.Run.id (Munin.Session.run session))\n"
},
{
"path": "packages/munin/lib/artifact.ml",
"content": "type kind = [ `dataset | `model | `checkpoint | `file | `dir | `other ]\ntype payload = [ `file | `dir ]\n\ntype t = {\n root : string;\n name : string;\n kind : kind;\n payload : payload;\n version : string;\n digest : string;\n materialized_rel_path : string;\n size_bytes : int;\n metadata : (string * Jsont.json) list;\n aliases : string list;\n producer_run_id : string option;\n consumer_run_ids : string list;\n created_at : float;\n}\n\nlet schema_version = 2\nlet name t = t.name\nlet kind t = t.kind\nlet payload t = t.payload\nlet version t = t.version\nlet digest t = t.digest\nlet size_bytes t = t.size_bytes\nlet metadata t = List.map (fun (k, v) -> (k, Value.of_json v)) t.metadata\nlet aliases t = t.aliases\nlet producer_run_id t = t.producer_run_id\nlet consumer_run_ids t = t.consumer_run_ids\nlet created_at t = t.created_at\nlet path t = Filename.concat t.root t.materialized_rel_path\nlet has_alias t alias = List.exists (String.equal alias) t.aliases\n\nlet kind_of_string = function\n | \"dataset\" -> `dataset\n | \"model\" -> `model\n | \"checkpoint\" -> `checkpoint\n | \"file\" -> `file\n | \"dir\" -> `dir\n | _ -> `other\n\nlet kind_to_string : kind -> string = function\n | `dataset -> \"dataset\"\n | `model -> \"model\"\n | `checkpoint -> \"checkpoint\"\n | `file -> \"file\"\n | `dir -> \"dir\"\n | `other -> \"other\"\n\nlet payload_of_string = function \"dir\" -> `dir | _ -> `file\n\nlet payload_to_string : payload -> string = function\n | `file -> \"file\"\n | `dir -> \"dir\"\n\nlet versions_dir root name =\n Filename.concat\n (Filename.concat (Filename.concat root \"artifacts\") name)\n \"versions\"\n\nlet manifest_path root name version =\n Filename.concat\n (Filename.concat (versions_dir root name) version)\n \"manifest.json\"\n\nlet load_manifest root path =\n try\n let json = Fs.read_file path |> Json_utils.json_of_string in\n let schema_ok =\n match\n Json_utils.json_mem \"schema_version\" json |> Json_utils.json_number\n with\n | Some v -> int_of_float v = schema_version\n | None -> false\n in\n if not schema_ok then None\n else\n match\n ( Json_utils.json_mem \"name\" json |> Json_utils.json_string,\n Json_utils.json_mem \"version\" json |> Json_utils.json_string,\n Json_utils.json_mem \"kind\" json |> Json_utils.json_string,\n Json_utils.json_mem \"payload\" json |> Json_utils.json_string,\n Json_utils.json_mem \"digest\" json |> Json_utils.json_string,\n Json_utils.json_mem \"path\" json |> Json_utils.json_string,\n Json_utils.json_mem \"size_bytes\" json |> Json_utils.json_number )\n with\n | ( Some name,\n Some version,\n Some kind,\n Some payload,\n Some digest,\n Some materialized_rel_path,\n Some size_bytes ) ->\n Some\n {\n root;\n name;\n kind = kind_of_string kind;\n payload = payload_of_string payload;\n version;\n digest;\n materialized_rel_path;\n size_bytes = int_of_float size_bytes;\n metadata =\n Json_utils.json_mem \"metadata\" json |> Json_utils.json_assoc;\n aliases =\n Json_utils.json_mem \"aliases\" json\n |> Json_utils.json_string_list;\n producer_run_id =\n Json_utils.json_mem \"producer_run_id\" json\n |> Json_utils.json_string;\n consumer_run_ids =\n Json_utils.json_mem \"consumer_run_ids\" json\n |> Json_utils.json_string_list;\n created_at =\n Option.value\n (Json_utils.json_mem \"created_at\" json\n |> Json_utils.json_number)\n ~default:0.0;\n }\n | _ -> None\n with _ -> None\n\nlet write_manifest root name version artifact =\n let json =\n Json_utils.json_obj\n ([\n (\"schema_version\", Jsont.Json.int schema_version);\n (\"name\", Jsont.Json.string artifact.name);\n (\"version\", Jsont.Json.string artifact.version);\n (\"kind\", Jsont.Json.string (kind_to_string artifact.kind));\n (\"payload\", Jsont.Json.string (payload_to_string artifact.payload));\n (\"digest\", Jsont.Json.string artifact.digest);\n (\"path\", Jsont.Json.string artifact.materialized_rel_path);\n (\"size_bytes\", Jsont.Json.int artifact.size_bytes);\n (\"metadata\", Json_utils.json_obj artifact.metadata);\n ( \"aliases\",\n Jsont.Json.list (List.map Jsont.Json.string artifact.aliases) );\n ( \"consumer_run_ids\",\n Jsont.Json.list\n (List.map Jsont.Json.string artifact.consumer_run_ids) );\n (\"created_at\", Jsont.Json.number artifact.created_at);\n ]\n @\n match artifact.producer_run_id with\n | None -> []\n | Some run_id -> [ (\"producer_run_id\", Jsont.Json.string run_id) ])\n in\n Fs.write_file\n (manifest_path root name version)\n (Json_utils.json_to_string ~pretty:true json ^ \"\\n\")\n\nlet version_number version =\n if String.length version >= 2 && version.[0] = 'v' then\n int_of_string_opt (String.sub version 1 (String.length version - 1))\n else None\n\nlet compare_version a b =\n match (version_number a.version, version_number b.version) with\n | Some a_num, Some b_num -> Int.compare a_num b_num\n | _ ->\n let by_name = String.compare a.name b.name in\n if by_name <> 0 then by_name else String.compare a.version b.version\n\nlet resolve_alias root name alias =\n Fs.list_dirs (versions_dir root name)\n |> List.filter_map (fun version ->\n load_manifest root (manifest_path root name version))\n |> List.filter (fun artifact -> has_alias artifact alias)\n |> List.sort compare_version |> List.rev\n |> function\n | artifact :: _ -> Some artifact\n | [] -> None\n\nlet load ~root ~name ~version =\n let path = manifest_path root name version in\n if Sys.file_exists path then load_manifest root path\n else resolve_alias root name version\n\nlet list ~root ?name ?kind ?alias ?producer_run ?consumer_run () =\n let names =\n match name with\n | Some name -> [ name ]\n | None -> Fs.list_dirs (Filename.concat root \"artifacts\")\n in\n List.concat_map\n (fun name ->\n Fs.list_dirs (versions_dir root name)\n |> List.filter_map (fun version ->\n load_manifest root (manifest_path root name version)))\n names\n |> List.filter (fun artifact ->\n Option.fold ~none:true ~some:(fun k -> artifact.kind = k) kind\n && Option.fold ~none:true ~some:(has_alias artifact) alias\n && Option.fold ~none:true\n ~some:(fun run_id -> artifact.producer_run_id = Some run_id)\n producer_run\n && Option.fold ~none:true\n ~some:(fun run_id ->\n List.exists (String.equal run_id) artifact.consumer_run_ids)\n consumer_run)\n |> List.sort (fun a b ->\n let by_name = String.compare a.name b.name in\n if by_name <> 0 then by_name else compare_version a b)\n\nlet create ~root ~name ~kind ~payload ~digest ~path:size_path ~metadata ~aliases\n ~producer_run_id =\n let version =\n let max_version =\n Fs.list_dirs (versions_dir root name)\n |> List.filter_map version_number\n |> List.fold_left max 0\n in\n Printf.sprintf \"v%d\" (max_version + 1)\n in\n let artifact =\n {\n root;\n name;\n kind;\n payload;\n version;\n digest;\n materialized_rel_path = size_path;\n size_bytes = Fs.path_size (Filename.concat root size_path);\n metadata;\n aliases;\n producer_run_id;\n consumer_run_ids = [];\n created_at = Unix.gettimeofday ();\n }\n in\n write_manifest root name version artifact;\n artifact\n\nlet add_consumer ~root ~name ~version run_id =\n match load ~root ~name ~version with\n | None -> ()\n | Some artifact ->\n if List.exists (String.equal run_id) artifact.consumer_run_ids then ()\n else\n let artifact =\n {\n artifact with\n consumer_run_ids = artifact.consumer_run_ids @ [ run_id ];\n }\n in\n write_manifest root name version artifact\n"
},
{
"path": "packages/munin/lib/artifact.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Versioned local artifacts.\n\n Artifacts are immutable named payloads with versions, aliases, and lineage.\n*)\n\n(** {1:types Types} *)\n\ntype kind =\n [ `dataset (** Dataset. *)\n | `model (** Model weights. *)\n | `checkpoint (** Checkpoint. *)\n | `file (** Single file. *)\n | `dir (** Directory tree. *)\n | `other (** Unclassified artifact. *) ]\n(** The type for logical artifact kinds. *)\n\ntype payload =\n [ `file (** Single file payload. *) | `dir (** Directory tree payload. *) ]\n(** The type for materialized payload kinds. *)\n\ntype t\n(** The type for artifact handles. *)\n\n(** {1:loading Loading} *)\n\nval load : root:string -> name:string -> version:string -> t option\n(** [load ~root ~name ~version] is the artifact named [name] at [version], if\n present.\n\n If [version] does not match an explicit version, it is resolved as an alias.\n Returns [None] if neither matches. *)\n\nval list :\n root:string ->\n ?name:string ->\n ?kind:kind ->\n ?alias:string ->\n ?producer_run:string ->\n ?consumer_run:string ->\n unit ->\n t list\n(** [list ~root ()] is the artifacts stored under [root], filtered when the\n optional selectors are provided. Results are sorted by name, then by version\n number. *)\n\n(** {1:identity Identity} *)\n\nval name : t -> string\n(** [name t] is the logical artifact name. *)\n\nval kind : t -> kind\n(** [kind t] is the artifact's logical kind. *)\n\nval payload : t -> payload\n(** [payload t] is the materialized payload kind. *)\n\nval version : t -> string\n(** [version t] is the explicit version such as [\"v1\"]. *)\n\n(** {1:content Content} *)\n\nval digest : t -> string\n(** [digest t] is the content-addressed SHA-256 digest of the payload. *)\n\nval path : t -> string\n(** [path t] is the absolute path to the materialized payload in the blob store.\n*)\n\nval size_bytes : t -> int\n(** [size_bytes t] is the total byte size of the materialized payload. *)\n\nval metadata : t -> (string * Value.t) list\n(** [metadata t] is the artifact metadata. *)\n\nval aliases : t -> string list\n(** [aliases t] is the alias list attached to the version. *)\n\nval has_alias : t -> string -> bool\n(** [has_alias t alias] is [true] iff [alias] points at [t]. *)\n\n(** {1:lineage Lineage} *)\n\nval producer_run_id : t -> string option\n(** [producer_run_id t] is the producing run identifier, if known. *)\n\nval consumer_run_ids : t -> string list\n(** [consumer_run_ids t] is the list of consuming run identifiers. *)\n\nval created_at : t -> float\n(** [created_at t] is the artifact creation timestamp ([Unix.gettimeofday] at\n creation time). *)\n\n(**/**)\n\nval create :\n root:string ->\n name:string ->\n kind:kind ->\n payload:payload ->\n digest:string ->\n path:string ->\n metadata:(string * Jsont.json) list ->\n aliases:string list ->\n producer_run_id:string option ->\n t\n\nval add_consumer :\n root:string -> name:string -> version:string -> string -> unit\n\n(**/**)\n"
},
{
"path": "packages/munin/lib/dune",
"content": "(library\n (name munin)\n (public_name munin)\n (wrapped false)\n (private_modules env fs json_utils event_log index)\n (libraries unix jsont jsont.bytesrw sha))\n"
},
{
"path": "packages/munin/lib/env.ml",
"content": "let root () =\n match Sys.getenv_opt \"RAVEN_TRACKING_DIR\" with\n | Some dir -> dir\n | None ->\n let data_home =\n match Sys.getenv_opt \"XDG_DATA_HOME\" with\n | Some dir -> dir\n | None ->\n Filename.concat\n (Filename.concat (Sys.getenv \"HOME\") \".local\")\n \"share\"\n in\n Filename.concat (Filename.concat data_home \"raven\") \"munin\"\n"
},
{
"path": "packages/munin/lib/event_log.ml",
"content": "type summary_mode = [ `Min | `Max | `Mean | `Last | `None ]\ntype goal = [ `Minimize | `Maximize ]\ntype media_kind = [ `Image | `Audio | `Table | `File ]\n\ntype event =\n | Metric of { step : int; timestamp : float; key : string; value : float }\n | Define_metric of {\n key : string;\n summary : summary_mode;\n step_metric : string option;\n goal : goal option;\n }\n | Media of {\n step : int;\n timestamp : float;\n key : string;\n kind : media_kind;\n path : string;\n }\n | Summary of (string * Value.t) list\n | Notes of string option\n | Tags of string list\n | Artifact_output of { name : string; version : string }\n | Artifact_input of { name : string; version : string }\n | Resumed of { at : float }\n | Finished of { status : string; ended_at : float }\n\nlet json_of_optional_string = function\n | Some value -> Jsont.Json.string value\n | None -> Json_utils.null\n\nlet optional_string_of_json json =\n match json with\n | Jsont.Null _ -> Some None\n | Jsont.String (value, _) -> Some (Some value)\n | _ -> None\n\nlet summary_mode_to_string = function\n | `Min -> \"min\"\n | `Max -> \"max\"\n | `Mean -> \"mean\"\n | `Last -> \"last\"\n | `None -> \"none\"\n\nlet summary_mode_of_string = function\n | \"min\" -> Some `Min\n | \"max\" -> Some `Max\n | \"mean\" -> Some `Mean\n | \"last\" -> Some `Last\n | \"none\" -> Some `None\n | _ -> None\n\nlet goal_to_string = function\n | `Minimize -> \"minimize\"\n | `Maximize -> \"maximize\"\n\nlet goal_of_string = function\n | \"minimize\" -> Some `Minimize\n | \"maximize\" -> Some `Maximize\n | _ -> None\n\nlet media_kind_to_string = function\n | `Image -> \"image\"\n | `Audio -> \"audio\"\n | `Table -> \"table\"\n | `File -> \"file\"\n\nlet media_kind_of_string = function\n | \"image\" -> Some `Image\n | \"audio\" -> Some `Audio\n | \"table\" -> Some `Table\n | \"file\" -> Some `File\n | _ -> None\n\nlet of_json json =\n match Json_utils.json_mem \"type\" json |> Json_utils.json_string with\n | Some \"metric\" -> (\n match\n ( Json_utils.json_mem \"step\" json |> Json_utils.json_number,\n Json_utils.json_mem \"timestamp\" json |> Json_utils.json_number,\n Json_utils.json_mem \"key\" json |> Json_utils.json_string,\n Json_utils.json_mem \"value\" json |> Json_utils.json_number )\n with\n | Some step, Some timestamp, Some key, Some value ->\n Some (Metric { step = int_of_float step; timestamp; key; value })\n | _ -> None)\n | Some \"define_metric\" -> (\n match\n ( Json_utils.json_mem \"key\" json |> Json_utils.json_string,\n Json_utils.json_mem \"summary\" json\n |> Json_utils.json_string\n |> Fun.flip Option.bind summary_mode_of_string )\n with\n | Some key, Some summary ->\n let step_metric =\n Json_utils.json_mem \"step_metric\" json |> Json_utils.json_string\n in\n let goal =\n Json_utils.json_mem \"goal\" json\n |> Json_utils.json_string\n |> Fun.flip Option.bind goal_of_string\n in\n Some (Define_metric { key; summary; step_metric; goal })\n | _ -> None)\n | Some \"media\" -> (\n match\n ( Json_utils.json_mem \"step\" json |> Json_utils.json_number,\n Json_utils.json_mem \"ts\" json |> Json_utils.json_number,\n Json_utils.json_mem \"key\" json |> Json_utils.json_string,\n Json_utils.json_mem \"kind\" json\n |> Json_utils.json_string\n |> Fun.flip Option.bind media_kind_of_string,\n Json_utils.json_mem \"path\" json |> Json_utils.json_string )\n with\n | Some step, Some ts, Some key, Some kind, Some path ->\n Some\n (Media { step = int_of_float step; timestamp = ts; key; kind; path })\n | _ -> None)\n | Some \"summary\" ->\n Some\n (Summary\n (Json_utils.json_mem \"values\" json\n |> Json_utils.json_assoc\n |> List.map (fun (k, v) -> (k, Value.of_json v))))\n | Some \"notes\" ->\n Json_utils.json_mem \"value\" json\n |> optional_string_of_json\n |> Option.map (fun value -> Notes value)\n | Some \"tags\" ->\n Some\n (Tags (Json_utils.json_mem \"values\" json |> Json_utils.json_string_list))\n | Some \"artifact_output\" -> (\n match\n ( Json_utils.json_mem \"name\" json |> Json_utils.json_string,\n Json_utils.json_mem \"version\" json |> Json_utils.json_string )\n with\n | Some name, Some version -> Some (Artifact_output { name; version })\n | _ -> None)\n | Some \"artifact_input\" -> (\n match\n ( Json_utils.json_mem \"name\" json |> Json_utils.json_string,\n Json_utils.json_mem \"version\" json |> Json_utils.json_string )\n with\n | Some name, Some version -> Some (Artifact_input { name; version })\n | _ -> None)\n | Some \"resumed\" -> (\n match Json_utils.json_mem \"at\" json |> Json_utils.json_number with\n | Some at -> Some (Resumed { at })\n | None -> None)\n | Some \"finished\" -> (\n match\n ( Json_utils.json_mem \"status\" json |> Json_utils.json_string,\n Json_utils.json_mem \"ended_at\" json |> Json_utils.json_number )\n with\n | Some status, Some ended_at -> Some (Finished { status; ended_at })\n | _ -> None)\n | _ -> None\n\nlet decode_line line =\n try Json_utils.json_of_string line |> of_json with _ -> None\n\nlet to_json = function\n | Metric { step; timestamp; key; value } ->\n Json_utils.json_obj\n [\n (\"type\", Jsont.Json.string \"metric\");\n (\"step\", Jsont.Json.int step);\n (\"timestamp\", Jsont.Json.number timestamp);\n (\"key\", Jsont.Json.string key);\n (\"value\", Jsont.Json.number value);\n ]\n | Define_metric { key; summary; step_metric; goal } ->\n Json_utils.json_obj\n ([\n (\"type\", Jsont.Json.string \"define_metric\");\n (\"key\", Jsont.Json.string key);\n (\"summary\", Jsont.Json.string (summary_mode_to_string summary));\n ]\n @ (match step_metric with\n | Some sm -> [ (\"step_metric\", Jsont.Json.string sm) ]\n | None -> [])\n @\n match goal with\n | Some g -> [ (\"goal\", Jsont.Json.string (goal_to_string g)) ]\n | None -> [])\n | Media { step; timestamp; key; kind; path } ->\n Json_utils.json_obj\n [\n (\"type\", Jsont.Json.string \"media\");\n (\"step\", Jsont.Json.int step);\n (\"ts\", Jsont.Json.number timestamp);\n (\"key\", Jsont.Json.string key);\n (\"kind\", Jsont.Json.string (media_kind_to_string kind));\n (\"path\", Jsont.Json.string path);\n ]\n | Summary values ->\n Json_utils.json_obj\n [\n (\"type\", Jsont.Json.string \"summary\");\n ( \"values\",\n Json_utils.json_obj\n (List.map (fun (k, v) -> (k, Value.to_json v)) values) );\n ]\n | Notes value ->\n Json_utils.json_obj\n [\n (\"type\", Jsont.Json.string \"notes\");\n (\"value\", json_of_optional_string value);\n ]\n | Tags values ->\n Json_utils.json_obj\n [\n (\"type\", Jsont.Json.string \"tags\");\n (\"values\", Jsont.Json.list (List.map Jsont.Json.string values));\n ]\n | Artifact_output { name; version } ->\n Json_utils.json_obj\n [\n (\"type\", Jsont.Json.string \"artifact_output\");\n (\"name\", Jsont.Json.string name);\n (\"version\", Jsont.Json.string version);\n ]\n | Artifact_input { name; version } ->\n Json_utils.json_obj\n [\n (\"type\", Jsont.Json.string \"artifact_input\");\n (\"name\", Jsont.Json.string name);\n (\"version\", Jsont.Json.string version);\n ]\n | Resumed { at } ->\n Json_utils.json_obj\n [ (\"type\", Jsont.Json.string \"resumed\"); (\"at\", Jsont.Json.number at) ]\n | Finished { status; ended_at } ->\n Json_utils.json_obj\n [\n (\"type\", Jsont.Json.string \"finished\");\n (\"status\", Jsont.Json.string status);\n (\"ended_at\", Jsont.Json.number ended_at);\n ]\n\nlet encode event = Json_utils.json_to_string (to_json event)\n\nlet read path =\n if not (Sys.file_exists path) then []\n else\n let ic = open_in path in\n let rec loop acc =\n match input_line ic with\n | line ->\n let acc =\n match decode_line line with\n | Some event -> event :: acc\n | None -> acc\n in\n loop acc\n | exception End_of_file -> List.rev acc\n in\n Fun.protect ~finally:(fun () -> close_in ic) (fun () -> loop [])\n"
},
{
"path": "packages/munin/lib/fs.ml",
"content": "let is_directory path =\n try (Unix.stat path).Unix.st_kind = Unix.S_DIR\n with Unix.Unix_error _ -> false\n\nlet ensure_dir path =\n let rec loop current =\n if current = \"\" || current = Filename.dir_sep then ()\n else if Sys.file_exists current then ()\n else (\n loop (Filename.dirname current);\n Unix.mkdir current 0o755)\n in\n loop path\n\nlet list_entries path =\n if Sys.file_exists path && is_directory path then\n Sys.readdir path |> Array.to_list |> List.sort String.compare\n else []\n\nlet list_dirs path =\n List.filter\n (fun name -> is_directory (Filename.concat path name))\n (list_entries path)\n\nlet read_file path =\n let ic = open_in_bin path in\n Fun.protect\n ~finally:(fun () -> close_in ic)\n (fun () -> really_input_string ic (in_channel_length ic))\n\nlet write_file path text =\n ensure_dir (Filename.dirname path);\n let oc = open_out_bin path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc text)\n\nlet append_line path line =\n ensure_dir (Filename.dirname path);\n let oc = open_out_gen [ Open_creat; Open_append; Open_binary ] 0o644 path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () ->\n output_string oc line;\n output_char oc '\\n')\n\nlet copy_file src dst =\n ensure_dir (Filename.dirname dst);\n let ic = open_in_bin src in\n let oc = open_out_bin dst in\n let buffer = Bytes.create 65536 in\n Fun.protect\n ~finally:(fun () ->\n close_in ic;\n close_out oc)\n (fun () ->\n let rec loop () =\n let count = input ic buffer 0 (Bytes.length buffer) in\n if count > 0 then (\n output oc buffer 0 count;\n loop ())\n in\n loop ())\n\nlet rec copy_tree src dst =\n if is_directory src then (\n ensure_dir dst;\n List.iter\n (fun name ->\n copy_tree (Filename.concat src name) (Filename.concat dst name))\n (list_entries src))\n else copy_file src dst\n\nlet rec remove_tree path =\n if is_directory path then (\n List.iter\n (fun name -> remove_tree (Filename.concat path name))\n (list_entries path);\n Unix.rmdir path)\n else Sys.remove path\n\nlet rec iter_tree ?(rel = \"\") root f =\n let path = if rel = \"\" then root else Filename.concat root rel in\n if is_directory path then (\n f rel `Dir;\n List.iter\n (fun name ->\n let child = if rel = \"\" then name else Filename.concat rel name in\n iter_tree ~rel:child root f)\n (list_entries path))\n else f rel `File\n\nlet sha256_file path = Sha256.to_hex (Sha256.file_fast path)\n\nlet sha256_path path =\n if is_directory path then (\n let ctx = Sha256.init () in\n iter_tree path (fun rel kind ->\n if rel <> \"\" then\n match kind with\n | `Dir ->\n Sha256.update_string ctx \"dir:\";\n Sha256.update_string ctx rel;\n Sha256.update_string ctx \"\\n\"\n | `File ->\n Sha256.update_string ctx \"file:\";\n Sha256.update_string ctx rel;\n Sha256.update_string ctx \":\";\n Sha256.update_string ctx (sha256_file (Filename.concat path rel));\n Sha256.update_string ctx \"\\n\");\n Sha256.to_hex (Sha256.finalize ctx))\n else sha256_file path\n\nlet file_size path =\n try (Unix.stat path).Unix.st_size with Unix.Unix_error _ -> 0\n\nlet rec path_size path =\n if is_directory path then\n list_entries path\n |> List.fold_left\n (fun acc name -> acc + path_size (Filename.concat path name))\n 0\n else file_size path\n\nlet command_output command =\n let ic = Unix.open_process_in (command ^ \" 2>/dev/null\") in\n let output =\n Fun.protect\n ~finally:(fun () -> ignore (Unix.close_process_in ic))\n (fun () ->\n let rec loop acc =\n match input_line ic with\n | line -> loop (line :: acc)\n | exception End_of_file -> List.rev acc\n in\n String.concat \"\\n\" (loop []))\n in\n if output = \"\" then None else Some output\n"
},
{
"path": "packages/munin/lib/index.ml",
"content": "type status = [ `running | `finished | `failed | `killed ]\n\ntype entry = {\n experiment : string;\n name : string option;\n group : string option;\n parent_id : string option;\n status : status;\n tags : string list;\n started_at : float;\n}\n\nlet index_path root = Filename.concat root \"index.json\"\n\nlet status_to_string = function\n | `running -> \"running\"\n | `finished -> \"finished\"\n | `failed -> \"failed\"\n | `killed -> \"killed\"\n\nlet entry_to_json entry =\n Json_utils.json_obj\n ([\n (\"experiment\", Jsont.Json.string entry.experiment);\n (\"status\", Jsont.Json.string (status_to_string entry.status));\n (\"tags\", Jsont.Json.list (List.map Jsont.Json.string entry.tags));\n (\"started_at\", Jsont.Json.number entry.started_at);\n ]\n @ (match entry.name with\n | Some n -> [ (\"name\", Jsont.Json.string n) ]\n | None -> [])\n @ (match entry.group with\n | Some g -> [ (\"group\", Jsont.Json.string g) ]\n | None -> [])\n @\n match entry.parent_id with\n | Some p -> [ (\"parent_id\", Jsont.Json.string p) ]\n | None -> [])\n\nlet entry_of_json json =\n match Json_utils.json_mem \"experiment\" json |> Json_utils.json_string with\n | None -> None\n | Some experiment ->\n let status : status =\n match Json_utils.json_mem \"status\" json |> Json_utils.json_string with\n | Some \"finished\" -> `finished\n | Some \"failed\" -> `failed\n | Some \"killed\" -> `killed\n | _ -> `running\n in\n let tags =\n Json_utils.json_mem \"tags\" json |> Json_utils.json_string_list\n in\n let started_at =\n Option.value\n (Json_utils.json_mem \"started_at\" json |> Json_utils.json_number)\n ~default:0.0\n in\n Some\n {\n experiment;\n name = Json_utils.json_mem \"name\" json |> Json_utils.json_string;\n group = Json_utils.json_mem \"group\" json |> Json_utils.json_string;\n parent_id =\n Json_utils.json_mem \"parent_id\" json |> Json_utils.json_string;\n status;\n tags;\n started_at;\n }\n\nlet read root =\n let path = index_path root in\n if not (Sys.file_exists path) then None\n else\n try\n let json = Fs.read_file path |> Json_utils.json_of_string in\n let tbl = Hashtbl.create 64 in\n List.iter\n (fun (id, value) ->\n match entry_of_json value with\n | Some entry -> Hashtbl.replace tbl id entry\n | None -> ())\n (Json_utils.json_assoc json);\n Some tbl\n with _ -> None\n\nlet write root tbl =\n let entries =\n Hashtbl.to_seq tbl |> List.of_seq\n |> List.sort (fun (a, _) (b, _) -> String.compare b a)\n |> List.map (fun (id, entry) -> (id, entry_to_json entry))\n in\n let json = Json_utils.json_obj entries in\n Fs.write_file (index_path root)\n (Json_utils.json_to_string ~pretty:true json ^ \"\\n\")\n\nlet modify root f =\n let tbl =\n match read root with Some tbl -> tbl | None -> Hashtbl.create 16\n in\n f tbl;\n write root tbl\n\nlet add root ~id entry = modify root (fun tbl -> Hashtbl.replace tbl id entry)\n\nlet update_status root ~id status =\n modify root (fun tbl ->\n match Hashtbl.find_opt tbl id with\n | Some entry -> Hashtbl.replace tbl id { entry with status }\n | None -> ())\n\nlet remove root ~id = modify root (fun tbl -> Hashtbl.remove tbl id)\n"
},
{
"path": "packages/munin/lib/json_utils.ml",
"content": "let null = Jsont.Null ((), Jsont.Meta.none)\n\nlet json_obj pairs =\n Jsont.Json.object'\n (List.map (fun (key, value) -> (Jsont.Json.name key, value)) pairs)\n\nlet json_mem name = function\n | Jsont.Object (members, _) -> (\n match Jsont.Json.find_mem name members with\n | Some (_, value) -> value\n | None -> null)\n | _ -> null\n\nlet json_string = function Jsont.String (value, _) -> Some value | _ -> None\nlet json_number = function Jsont.Number (value, _) -> Some value | _ -> None\nlet json_bool = function Jsont.Bool (value, _) -> Some value | _ -> None\n\nlet json_string_list = function\n | Jsont.Array (items, _) ->\n List.filter_map\n (function Jsont.String (value, _) -> Some value | _ -> None)\n items\n | _ -> []\n\nlet json_assoc = function\n | Jsont.Object (members, _) ->\n List.map (fun ((key, _), value) -> (key, value)) members\n | _ -> []\n\nlet json_of_string text =\n match Jsont_bytesrw.decode_string Jsont.json text with\n | Ok json -> json\n | Error message -> failwith message\n\nlet json_to_string ?(pretty = false) json =\n let format = if pretty then Jsont.Indent else Jsont.Minify in\n match Jsont_bytesrw.encode_string ~format Jsont.json json with\n | Ok text -> text\n | Error message -> failwith message\n"
},
{
"path": "packages/munin/lib/munin.ml",
"content": "module Value = Value\nmodule Artifact = Artifact\nmodule Run = Run\nmodule Run_monitor = Run_monitor\nmodule Session = Session\nmodule Store = Store\n"
},
{
"path": "packages/munin/lib/munin.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Local experiment tracking for Raven.\n\n Munin is a local-first experiment tracker. Start with {!Session} to write\n runs, {!Run} to read them back, {!Run_monitor} for live polling, {!Store}\n for discovery, and {!Artifact} for versioned payloads.\n\n {1:library Library [munin]}\n {!modules:Value Session Run Run_monitor Store Artifact} *)\n\nmodule Value = Value\nmodule Artifact = Artifact\nmodule Run = Run\nmodule Run_monitor = Run_monitor\nmodule Session = Session\nmodule Store = Store\n"
},
{
"path": "packages/munin/lib/run.ml",
"content": "type status = [ `running | `finished | `failed | `killed ]\ntype metric = { step : int; timestamp : float; value : float }\n\ntype provenance = {\n notes : string option;\n command : string list;\n cwd : string;\n hostname : string option;\n pid : int;\n git_commit : string option;\n git_dirty : bool option;\n env : (string * string) list;\n}\n\ntype metric_def = {\n summary : [ `Min | `Max | `Mean | `Last | `None ];\n step_metric : string option;\n goal : [ `Minimize | `Maximize ] option;\n}\n\ntype media_entry = {\n step : int;\n timestamp : float;\n kind : [ `Image | `Audio | `Table | `File ];\n path : string;\n}\n\n(* Heavy data loaded on demand from manifest + events *)\ntype full = {\n params : (string * Value.t) list;\n provenance : provenance;\n ended_at : float option;\n summary : (string * Value.t) list;\n latest_metrics : (string * metric) list;\n histories : (string * metric list) list;\n metric_defs : (string * metric_def) list;\n media : (string * media_entry list) list;\n input_artifacts : Artifact.t list;\n output_artifacts : Artifact.t list;\n}\n\n(* Header fields are always available without I/O *)\ntype t = {\n root : string;\n id : string;\n dir : string;\n experiment_name : string;\n name : string option;\n group : string option;\n parent_id : string option;\n started_at : float;\n status : status;\n tags : string list;\n full : full Lazy.t;\n}\n\nlet schema_version = 2\n\n(* Header accessors — no I/O *)\n\nlet id t = t.id\nlet dir t = t.dir\nlet experiment_name t = t.experiment_name\nlet name t = t.name\nlet group t = t.group\nlet parent_id t = t.parent_id\nlet started_at t = t.started_at\nlet status t = t.status\nlet tags t = t.tags\nlet resumable t = t.status = `running\n\n(* Full accessors — forces lazy on first access *)\n\nlet full t = Lazy.force t.full\nlet params t = (full t).params\nlet provenance t = (full t).provenance\nlet notes t = (full t).provenance.notes\nlet ended_at t = (full t).ended_at\nlet summary t = (full t).summary\nlet find_param t key = List.assoc_opt key (full t).params\nlet find_summary t key = List.assoc_opt key (full t).summary\nlet latest_metrics t = (full t).latest_metrics\nlet metric_keys t = List.map fst (full t).latest_metrics\nlet input_artifacts t = (full t).input_artifacts\nlet output_artifacts t = (full t).output_artifacts\nlet metric_defs t = (full t).metric_defs\nlet media_keys t = List.map fst (full t).media\n\nlet media_history t key =\n match List.assoc_opt key (full t).media with\n | Some entries -> entries\n | None -> []\n\nlet metric_history t key =\n match List.assoc_opt key (full t).histories with\n | Some history -> history\n | None -> []\n\n(* Paths *)\n\nlet run_dir ~root ~experiment id =\n Filename.concat\n (Filename.concat\n (Filename.concat (Filename.concat root \"experiments\") experiment)\n \"runs\")\n id\n\nlet manifest_path root experiment id =\n Filename.concat (run_dir ~root ~experiment id) \"run.json\"\n\nlet events_path dir = Filename.concat dir \"events.jsonl\"\n\n(* Parsing helpers *)\n\nlet status_of_string = function\n | \"finished\" -> `finished\n | \"failed\" -> `failed\n | \"killed\" -> `killed\n | _ -> `running\n\nlet push_tag seen acc tag =\n if Hashtbl.mem seen tag then acc\n else (\n Hashtbl.replace seen tag ();\n tag :: acc)\n\nlet provenance_of_json json =\n let env_json = Json_utils.json_mem \"env\" json in\n {\n notes = Json_utils.json_mem \"notes\" json |> Json_utils.json_string;\n command = Json_utils.json_mem \"command\" json |> Json_utils.json_string_list;\n cwd =\n Option.value\n (Json_utils.json_mem \"cwd\" json |> Json_utils.json_string)\n ~default:\"\";\n hostname = Json_utils.json_mem \"hostname\" json |> Json_utils.json_string;\n pid =\n Option.value\n (Json_utils.json_mem \"pid\" json |> Json_utils.json_number)\n ~default:0.0\n |> int_of_float;\n git_commit = Json_utils.json_mem \"git_commit\" json |> Json_utils.json_string;\n git_dirty = Json_utils.json_mem \"git_dirty\" json |> Json_utils.json_bool;\n env =\n Json_utils.json_assoc env_json\n |> List.filter_map (fun (key, value) ->\n Json_utils.json_string value |> Option.map (fun text -> (key, text)));\n }\n\nlet sorted_of_hashtbl tbl =\n Hashtbl.to_seq tbl |> List.of_seq\n |> List.sort (fun (a, _) (b, _) -> String.compare a b)\n\n(* Materialize full data from manifest JSON + event log *)\nlet materialize_full root dir manifest_json =\n let tag_seen = Hashtbl.create 8 in\n let initial_tags =\n Json_utils.json_mem \"tags\" manifest_json\n |> Json_utils.json_string_list\n |> List.fold_left (push_tag tag_seen) []\n |> List.rev\n in\n let params =\n Json_utils.json_mem \"params\" manifest_json\n |> Json_utils.json_assoc\n |> List.map (fun (k, v) -> (k, Value.of_json v))\n in\n let summary_table = Hashtbl.create 8 in\n let history_table = Hashtbl.create 16 in\n let latest_table = Hashtbl.create 16 in\n let metric_def_table = Hashtbl.create 8 in\n let media_table = Hashtbl.create 8 in\n let input_seen = Hashtbl.create 8 in\n let output_seen = Hashtbl.create 8 in\n let input_artifacts = ref [] in\n let output_artifacts = ref [] in\n let tags = ref initial_tags in\n let status = ref `running in\n let ended_at = ref None in\n let notes =\n ref\n (Json_utils.json_mem \"provenance\" manifest_json |> provenance_of_json)\n .notes\n in\n List.iter\n (function\n | Event_log.Metric { step; timestamp; key; value } ->\n let metric = { step; timestamp; value } in\n let history =\n match Hashtbl.find_opt history_table key with\n | Some history -> history\n | None -> []\n in\n Hashtbl.replace history_table key (metric :: history);\n Hashtbl.replace latest_table key metric\n | Define_metric { key; summary; step_metric; goal } ->\n Hashtbl.replace metric_def_table key { summary; step_metric; goal }\n | Media { step; timestamp; key; kind; path } ->\n let abs_path = Filename.concat dir path in\n let entry = { step; timestamp; kind; path = abs_path } in\n let prev =\n match Hashtbl.find_opt media_table key with\n | Some l -> l\n | None -> []\n in\n Hashtbl.replace media_table key (entry :: prev)\n | Summary values ->\n List.iter\n (fun (key, value) -> Hashtbl.replace summary_table key value)\n values\n | Notes value -> notes := value\n | Tags values -> tags := List.fold_left (push_tag tag_seen) !tags values\n | Artifact_output { name; version } ->\n let key = name ^ \":\" ^ version in\n if not (Hashtbl.mem output_seen key) then (\n Hashtbl.replace output_seen key ();\n match Artifact.load ~root ~name ~version with\n | Some artifact -> output_artifacts := artifact :: !output_artifacts\n | None -> ())\n | Artifact_input { name; version } ->\n let key = name ^ \":\" ^ version in\n if not (Hashtbl.mem input_seen key) then (\n Hashtbl.replace input_seen key ();\n match Artifact.load ~root ~name ~version with\n | Some artifact -> input_artifacts := artifact :: !input_artifacts\n | None -> ())\n | Resumed _ ->\n ended_at := None;\n status := `running\n | Finished { status = status_string; ended_at = finished_at } ->\n status := status_of_string status_string;\n ended_at := Some finished_at)\n (Event_log.read (events_path dir));\n let latest_metrics = sorted_of_hashtbl latest_table in\n let histories =\n sorted_of_hashtbl history_table\n |> List.map (fun (key, history) -> (key, List.rev history))\n in\n let metric_defs = sorted_of_hashtbl metric_def_table in\n (* Auto-compute summaries from define_metric declarations. Explicit\n set_summary always wins; auto-summary only fills gaps. *)\n Hashtbl.iter\n (fun key (def : metric_def) ->\n if not (Hashtbl.mem summary_table key) then\n match Hashtbl.find_opt history_table key with\n | None | Some [] -> ()\n | Some history ->\n let auto =\n match def.summary with\n | `Min ->\n Some\n (List.fold_left\n (fun acc (m : metric) -> Float.min acc m.value)\n Float.infinity history)\n | `Max ->\n Some\n (List.fold_left\n (fun acc (m : metric) -> Float.max acc m.value)\n Float.neg_infinity history)\n | `Mean ->\n let n = List.length history in\n let sum =\n List.fold_left\n (fun acc (m : metric) -> acc +. m.value)\n 0. history\n in\n Some (sum /. Float.of_int n)\n | `Last -> Some (List.hd history).value\n | `None -> None\n in\n Option.iter\n (fun v -> Hashtbl.replace summary_table key (`Float v))\n auto)\n metric_def_table;\n let summary = sorted_of_hashtbl summary_table in\n let media =\n sorted_of_hashtbl media_table\n |> List.map (fun (key, entries) -> (key, List.rev entries))\n in\n let base_provenance =\n Json_utils.json_mem \"provenance\" manifest_json |> provenance_of_json\n in\n ( !status,\n List.rev !tags,\n {\n params;\n provenance = { base_provenance with notes = !notes };\n ended_at = !ended_at;\n summary;\n latest_metrics;\n histories;\n metric_defs;\n media;\n input_artifacts = List.rev !input_artifacts;\n output_artifacts = List.rev !output_artifacts;\n } )\n\n(* Full eager load — reads manifest + events immediately *)\nlet load ~root ~experiment ~id =\n let path = manifest_path root experiment id in\n if not (Sys.file_exists path) then None\n else\n try\n let json = Fs.read_file path |> Json_utils.json_of_string in\n let schema_ok =\n match\n Json_utils.json_mem \"schema_version\" json |> Json_utils.json_number\n with\n | Some value -> int_of_float value = schema_version\n | None -> false\n in\n if not schema_ok then None\n else\n let dir = run_dir ~root ~experiment id in\n let name = Json_utils.json_mem \"name\" json |> Json_utils.json_string in\n let group =\n Json_utils.json_mem \"group\" json |> Json_utils.json_string\n in\n let parent_id =\n Json_utils.json_mem \"parent_id\" json |> Json_utils.json_string\n in\n let started_at =\n Option.value\n (Json_utils.json_mem \"started_at\" json |> Json_utils.json_number)\n ~default:0.0\n in\n let status, tags, full_data = materialize_full root dir json in\n Some\n {\n root;\n id;\n dir;\n experiment_name = experiment;\n name;\n group;\n parent_id;\n started_at;\n status;\n tags;\n full = Lazy.from_val full_data;\n }\n with _ -> None\n\n(* Lazy load from index — reads manifest + events only when full data\n accessed *)\nlet load_from_index ~root id (entry : Index.entry) =\n let dir = run_dir ~root ~experiment:entry.experiment id in\n let full =\n lazy\n (let path = manifest_path root entry.experiment id in\n let json = Fs.read_file path |> Json_utils.json_of_string in\n let _status, _tags, full_data = materialize_full root dir json in\n full_data)\n in\n {\n root;\n id;\n dir;\n experiment_name = entry.experiment;\n name = entry.name;\n group = entry.group;\n parent_id = entry.parent_id;\n started_at = entry.started_at;\n status = entry.status;\n tags = entry.tags;\n full;\n }\n\nlet list ~root ~experiment ?status:status_filter ?tag ?parent\n ?group:group_filter () =\n let runs_dir =\n Filename.concat\n (Filename.concat (Filename.concat root \"experiments\") experiment)\n \"runs\"\n in\n Fs.list_dirs runs_dir\n |> List.filter_map (fun id -> load ~root ~experiment ~id)\n |> List.filter (fun run ->\n Option.fold ~none:true ~some:(fun s -> status run = s) status_filter\n && Option.fold ~none:true\n ~some:(fun tag -> List.exists (String.equal tag) (tags run))\n tag\n && Option.fold ~none:true\n ~some:(fun parent -> parent_id run = Some parent)\n parent\n && Option.fold ~none:true ~some:(fun g -> group run = Some g) group_filter)\n |> List.sort (fun a b -> String.compare (id b) (id a))\n\nlet children t = list ~root:t.root ~experiment:t.experiment_name ~parent:t.id ()\n"
},
{
"path": "packages/munin/lib/run.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Persisted tracked runs.\n\n Runs are the durable tracked objects of Munin. They expose immutable\n manifest data together with materialized state rebuilt from the append-only\n event log. *)\n\n(** {1:types Types} *)\n\ntype status =\n [ `running (** Run is actively logging. *)\n | `finished (** Run completed successfully. *)\n | `failed (** Run terminated with an error. *)\n | `killed (** Run was manually terminated. *) ]\n(** The type for run status values. *)\n\ntype metric = {\n step : int; (** Step counter at which the sample was logged. *)\n timestamp : float; (** Wall-clock time of the sample. *)\n value : float; (** Scalar metric value. *)\n}\n(** The type for scalar metric observations. *)\n\ntype provenance = {\n notes : string option; (** Free-form run note. *)\n command : string list; (** Command line that started the run. *)\n cwd : string; (** Working directory at run start. *)\n hostname : string option; (** Machine hostname. *)\n pid : int; (** Process identifier. *)\n git_commit : string option; (** Git HEAD commit hash. *)\n git_dirty : bool option; (** Whether the working tree was dirty. *)\n env : (string * string) list; (** Captured environment variables. *)\n}\n(** The type for run provenance. *)\n\ntype metric_def = {\n summary : [ `Min | `Max | `Mean | `Last | `None ];\n (** How the run summary value is computed from history. *)\n step_metric : string option;\n (** Another metric to use as x-axis (e.g. [\"epoch\"]). *)\n goal : [ `Minimize | `Maximize ] option;\n (** Whether lower or higher values are better. *)\n}\n(** The type for metric definitions. Declares how a metric should be summarised\n and plotted. *)\n\ntype media_entry = {\n step : int; (** Step counter at which the media was logged. *)\n timestamp : float; (** Wall-clock time. *)\n kind : [ `Image | `Audio | `Table | `File ]; (** Media type for renderers. *)\n path : string; (** Absolute path to stored file. *)\n}\n(** The type for media log entries. *)\n\ntype t\n(** The type for run handles. *)\n\n(** {1:loading Loading} *)\n\nval load : root:string -> experiment:string -> id:string -> t option\n(** [load ~root ~experiment ~id] is the run [id] in [experiment], if present.\n Returns [None] if the manifest is missing, has an incompatible schema\n version, or cannot be read. *)\n\nval list :\n root:string ->\n experiment:string ->\n ?status:status ->\n ?tag:string ->\n ?parent:string ->\n ?group:string ->\n unit ->\n t list\n(** [list ~root ~experiment ()] is the runs persisted for [experiment], filtered\n when the optional selectors are provided. [parent] filters by parent run\n identifier. Results are sorted by identifier descending (newest first). *)\n\n(** {1:identity Identity} *)\n\nval id : t -> string\n(** [id t] is the unique run identifier. *)\n\nval dir : t -> string\n(** [dir t] is the absolute path to the run directory. *)\n\nval experiment_name : t -> string\n(** [experiment_name t] is the containing experiment name. *)\n\nval name : t -> string option\n(** [name t] is the optional human-readable run name. *)\n\nval group : t -> string option\n(** [group t] is the optional run group for flat grouping (e.g. sweeps). *)\n\nval parent_id : t -> string option\n(** [parent_id t] is the parent run identifier, if any. *)\n\n(** {1:status Status} *)\n\nval started_at : t -> float\n(** [started_at t] is the run start timestamp. *)\n\nval ended_at : t -> float option\n(** [ended_at t] is the run completion timestamp, if any. *)\n\nval status : t -> status\n(** [status t] is the current run status. *)\n\nval resumable : t -> bool\n(** [resumable t] is [true] iff [status t] is [`running]. *)\n\n(** {1:provenance Provenance} *)\n\nval provenance : t -> provenance\n(** [provenance t] is the run provenance. *)\n\nval notes : t -> string option\n(** [notes t] is the latest run note, if any. *)\n\n(** {1:metadata Metadata} *)\n\nval tags : t -> string list\n(** [tags t] is the run tag list. *)\n\nval params : t -> (string * Value.t) list\n(** [params t] is the immutable run parameter set. *)\n\nval find_param : t -> string -> Value.t option\n(** [find_param t key] is the parameter value for [key], if present. *)\n\nval summary : t -> (string * Value.t) list\n(** [summary t] is the run summary map, sorted alphabetically by key. Later\n writes replace earlier values. *)\n\nval find_summary : t -> string -> Value.t option\n(** [find_summary t key] is the summary value for [key], if present. *)\n\n(** {1:metrics Metrics} *)\n\nval metric_keys : t -> string list\n(** [metric_keys t] is the sorted list of metric keys observed in [t]. *)\n\nval latest_metrics : t -> (string * metric) list\n(** [latest_metrics t] is the latest scalar metric value per key, sorted\n alphabetically by key. *)\n\nval metric_history : t -> string -> metric list\n(** [metric_history t key] is the full history for [key] in chronological order.\n Returns the empty list if [key] has no samples. *)\n\nval metric_defs : t -> (string * metric_def) list\n(** [metric_defs t] is the metric definitions declared via\n {!Session.define_metric}, sorted alphabetically by key. *)\n\n(** {1:media Media} *)\n\nval media_keys : t -> string list\n(** [media_keys t] is the sorted list of media keys logged in [t]. *)\n\nval media_history : t -> string -> media_entry list\n(** [media_history t key] is the media entries for [key] in chronological order.\n Returns the empty list if [key] has no entries. *)\n\n(** {1:relations Relations} *)\n\nval children : t -> t list\n(** [children t] is the list of child runs of [t]. Performs a filesystem scan of\n the experiment directory. *)\n\nval input_artifacts : t -> Artifact.t list\n(** [input_artifacts t] is the list of artifacts consumed by [t]. *)\n\nval output_artifacts : t -> Artifact.t list\n(** [output_artifacts t] is the list of artifacts produced by [t]. *)\n\n(**/**)\n\nval status_of_string : string -> status\nval load_from_index : root:string -> string -> Index.entry -> t\n\n(**/**)\n"
},
{
"path": "packages/munin/lib/run_monitor.ml",
"content": "type live_status = [ `Live | `Stopped | `Done of Run.status ]\n\ntype t = {\n run : Run.t;\n mutable ic : in_channel option;\n mutable pos : int;\n mutable last_event_time : float;\n latest : (string, Run.metric) Hashtbl.t;\n histories : (string, (int * float) list) Hashtbl.t;\n defs : (string, Run.metric_def) Hashtbl.t;\n mutable finished : Run.status option;\n}\n\nlet stopped_timeout = 5.0\nlet events_path dir = Filename.concat dir \"events.jsonl\"\n\nlet start run =\n {\n run;\n ic = None;\n pos = 0;\n last_event_time = Unix.gettimeofday ();\n latest = Hashtbl.create 16;\n histories = Hashtbl.create 16;\n defs = Hashtbl.create 8;\n finished = None;\n }\n\nlet close t =\n Option.iter close_in t.ic;\n t.ic <- None\n\nlet read_new_events t =\n let path = events_path (Run.dir t.run) in\n if not (Sys.file_exists path) then []\n else\n let stat = Unix.stat path in\n let size = stat.Unix.st_size in\n if size <= t.pos then []\n else\n let ic =\n match t.ic with\n | Some ic ->\n (* Check for truncation/rotation *)\n if size < t.pos then (\n close_in ic;\n let new_ic = open_in path in\n t.ic <- Some new_ic;\n t.pos <- 0;\n new_ic)\n else ic\n | None ->\n let ic = open_in path in\n t.ic <- Some ic;\n ic\n in\n seek_in ic t.pos;\n let events = ref [] in\n (try\n while true do\n let line = input_line ic in\n match Event_log.decode_line line with\n | Some event -> events := event :: !events\n | None -> ()\n done\n with End_of_file -> ());\n t.pos <- pos_in ic;\n List.rev !events\n\nlet process_event t = function\n | Event_log.Metric { step; timestamp; key; value } ->\n let metric = Run.{ step; timestamp; value } in\n Hashtbl.replace t.latest key metric;\n let history =\n match Hashtbl.find_opt t.histories key with Some h -> h | None -> []\n in\n Hashtbl.replace t.histories key ((step, value) :: history);\n t.last_event_time <- timestamp\n | Define_metric { key; summary = s; step_metric; goal } ->\n Hashtbl.replace t.defs key { Run.summary = s; step_metric; goal }\n | Finished { status; ended_at = _ } ->\n t.finished <- Some (Run.status_of_string status)\n | Resumed _ ->\n t.finished <- None;\n t.last_event_time <- Unix.gettimeofday ()\n | Media _ | Summary _ | Notes _ | Tags _ | Artifact_output _\n | Artifact_input _ ->\n ()\n\nlet poll t =\n let events = read_new_events t in\n List.iter (process_event t) events\n\nlet live_status t =\n match t.finished with\n | Some status -> `Done status\n | None ->\n let elapsed = Unix.gettimeofday () -. t.last_event_time in\n if elapsed > stopped_timeout then `Stopped else `Live\n\nlet metrics t =\n Hashtbl.to_seq t.latest |> List.of_seq\n |> List.sort (fun (a, _) (b, _) -> String.compare a b)\n\nlet history t key =\n match Hashtbl.find_opt t.histories key with\n | Some h -> List.rev h\n | None -> []\n\nlet metric_defs t =\n Hashtbl.to_seq t.defs |> List.of_seq\n |> List.sort (fun (a, _) (b, _) -> String.compare a b)\n\nlet contains_sub ~sub s =\n let ls = String.length sub and lk = String.length s in\n ls <= lk\n &&\n let rec loop i = i <= lk - ls && (String.sub s i ls = sub || loop (i + 1)) in\n loop 0\n\nlet is_loss_like key =\n let key = String.lowercase_ascii key in\n contains_sub ~sub:\"loss\" key || contains_sub ~sub:\"error\" key\n\nlet best t key =\n match Hashtbl.find_opt t.histories key with\n | None | Some [] -> None\n | Some history ->\n let minimize =\n match Hashtbl.find_opt t.defs key with\n | Some { goal = Some `Minimize; _ } -> true\n | Some { goal = Some `Maximize; _ } -> false\n | _ -> is_loss_like key\n in\n let compare =\n if minimize then fun a b -> Float.compare a b\n else fun a b -> Float.compare b a\n in\n let best_step, best_value =\n List.fold_left\n (fun (bs, bv) (s, v) -> if compare v bv < 0 then (s, v) else (bs, bv))\n (List.hd history) (List.tl history)\n in\n let timestamp =\n match Hashtbl.find_opt t.latest key with\n | Some m -> m.timestamp\n | None -> 0.0\n in\n Some Run.{ step = best_step; timestamp; value = best_value }\n"
},
{
"path": "packages/munin/lib/run_monitor.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Incremental run monitoring.\n\n [Run_monitor] polls a run's event log and maintains aggregated metric state.\n Used by the TUI and future web dashboard for live updates without re-reading\n the entire log. *)\n\ntype t\n(** The type for run monitors. *)\n\ntype live_status =\n [ `Live (** Events are still arriving. *)\n | `Stopped (** No events for the timeout period (5 seconds). *)\n | `Done of Run.status (** [Finished] event received. *) ]\n(** The type for live run status. *)\n\nval start : Run.t -> t\n(** [start run] opens the run's event log for incremental reading. The event log\n file handle is acquired lazily on the first {!poll}. Must be paired with\n {!close}. *)\n\nval poll : t -> unit\n(** [poll t] reads new events since the last poll and updates state. *)\n\nval close : t -> unit\n(** [close t] releases the file handle. *)\n\n(** {1:state State} *)\n\nval live_status : t -> live_status\n(** [live_status t] is the current run status based on event activity. *)\n\nval metrics : t -> (string * Run.metric) list\n(** [metrics t] is the latest metric value per key, sorted alphabetically. *)\n\nval history : t -> string -> (int * float) list\n(** [history t key] is the [(step, value)] history for [key] in chronological\n order. Returns the empty list if [key] has no samples. *)\n\nval metric_defs : t -> (string * Run.metric_def) list\n(** [metric_defs t] is the metric definitions declared so far, sorted\n alphabetically by key. *)\n\nval best : t -> string -> Run.metric option\n(** [best t key] is the best observation for [key] according to\n {!Session.define_metric} goal, or a heuristic if undefined (keys containing\n \"loss\" or \"error\" prefer lower values). Returns [None] if no samples exist\n for [key]. *)\n"
},
{
"path": "packages/munin/lib/session.ml",
"content": "let err_not_resumable = \"Munin.Session.resume: run is not resumable\"\nlet err_missing_manifest = \"Munin.Session.run: missing run manifest\"\nlet err_closed_session = \"Munin.Session.log_artifact: closed session\"\nlet err_path_missing = \"Munin.Session.log_artifact: path does not exist: \"\nlet err_media_missing = \"Munin.Session.log_media: path does not exist: \"\n\ntype t = {\n root : string;\n experiment : string;\n id : string;\n dir : string;\n mutex : Mutex.t;\n mutable closed : bool;\n}\n\nlet schema_version = 2\nlet manifest_path dir = Filename.concat dir \"run.json\"\nlet events_path dir = Filename.concat dir \"events.jsonl\"\nlet random_state = lazy (Random.State.make_self_init ())\n\nlet generate_id () =\n let state = Lazy.force random_state in\n let now = Unix.gettimeofday () in\n let tm = Unix.localtime now in\n let stamp =\n Printf.sprintf \"%04d-%02d-%02d_%02d-%02d-%02d\" (tm.Unix.tm_year + 1900)\n (tm.Unix.tm_mon + 1) tm.Unix.tm_mday tm.Unix.tm_hour tm.Unix.tm_min\n tm.Unix.tm_sec\n in\n let suffix = Printf.sprintf \"%04x\" (Random.State.int state 0x10000) in\n stamp ^ \"_\" ^ suffix\n\nlet status_to_string = function\n | `finished -> \"finished\"\n | `failed -> \"failed\"\n | `killed -> \"killed\"\n\nlet git_output cwd args =\n let command =\n String.concat \" \" (List.map Filename.quote (\"git\" :: \"-C\" :: cwd :: args))\n in\n Fs.command_output command\n\nlet detect_git_commit cwd = git_output cwd [ \"rev-parse\"; \"HEAD\" ]\n\nlet detect_git_dirty cwd =\n match git_output cwd [ \"status\"; \"--porcelain\"; \"--untracked-files=no\" ] with\n | None -> None\n | Some output -> Some (output <> \"\")\n\nlet capture_env_vars names =\n List.filter_map\n (fun name -> Option.map (fun value -> (name, value)) (Sys.getenv_opt name))\n names\n\nlet first_some a b = match a with Some _ -> a | None -> b\n\nlet root_of_run_dir dir =\n Filename.dirname (Filename.dirname (Filename.dirname (Filename.dirname dir)))\n\nlet with_lock t f =\n Mutex.lock t.mutex;\n Fun.protect ~finally:(fun () -> Mutex.unlock t.mutex) f\n\nlet append_event t event =\n Fs.append_line (events_path t.dir) (Event_log.encode event)\n\nlet write_manifest path json =\n Fs.write_file path (Json_utils.json_to_string ~pretty:true json ^ \"\\n\")\n\nlet optional_field key f = function None -> [] | Some v -> [ (key, f v) ]\n\nlet provenance_json ?notes ~command ~cwd ~hostname ~pid ~git_commit ~git_dirty\n ~env () =\n Json_utils.json_obj\n ([\n (\"command\", Jsont.Json.list (List.map Jsont.Json.string command));\n (\"cwd\", Jsont.Json.string cwd);\n (\"pid\", Jsont.Json.int pid);\n ( \"env\",\n Json_utils.json_obj\n (List.map (fun (k, v) -> (k, Jsont.Json.string v)) env) );\n ]\n @ optional_field \"notes\" Jsont.Json.string notes\n @ optional_field \"hostname\" Jsont.Json.string hostname\n @ optional_field \"git_commit\" Jsont.Json.string git_commit\n @ optional_field \"git_dirty\" Jsont.Json.bool git_dirty)\n\nlet make_manifest ~id ~experiment ~started_at ?name ?group ?parent ~tags ~params\n ~provenance () =\n Json_utils.json_obj\n ([\n (\"schema_version\", Jsont.Json.int schema_version);\n (\"id\", Jsont.Json.string id);\n (\"experiment\", Jsont.Json.string experiment);\n (\"started_at\", Jsont.Json.number started_at);\n (\"tags\", Jsont.Json.list (List.map Jsont.Json.string tags));\n ( \"params\",\n Json_utils.json_obj\n (List.map (fun (k, v) -> (k, Value.to_json v)) params) );\n (\"provenance\", provenance);\n ]\n @ optional_field \"name\" Jsont.Json.string name\n @ optional_field \"group\" Jsont.Json.string group\n @ optional_field \"parent_id\"\n (fun run -> Jsont.Json.string (Run.id run))\n parent)\n\nlet start ?root ~experiment ?name ?group ?parent ?(tags = []) ?(params = [])\n ?notes ?(capture_env = []) ?command ?cwd ?hostname ?pid ?git_commit\n ?git_dirty ?env () =\n let root = Option.value root ~default:(Env.root ()) in\n Fs.ensure_dir (Filename.concat root \"experiments\");\n Fs.ensure_dir (Filename.concat root \"artifacts\");\n Fs.ensure_dir (Filename.concat (Filename.concat root \"blobs\") \"sha256\");\n let id = generate_id () in\n let dir =\n Filename.concat\n (Filename.concat\n (Filename.concat (Filename.concat root \"experiments\") experiment)\n \"runs\")\n id\n in\n Fs.ensure_dir dir;\n let cwd = Option.value cwd ~default:(Sys.getcwd ()) in\n let command = Option.value command ~default:(Array.to_list Sys.argv) in\n let hostname =\n match hostname with\n | Some hostname -> Some hostname\n | None -> Some (Unix.gethostname ())\n in\n let pid = Option.value pid ~default:(Unix.getpid ()) in\n let git_commit = first_some git_commit (detect_git_commit cwd) in\n let git_dirty = first_some git_dirty (detect_git_dirty cwd) in\n let env = Option.value env ~default:(capture_env_vars capture_env) in\n let provenance =\n provenance_json ?notes ~command ~cwd ~hostname ~pid ~git_commit ~git_dirty\n ~env ()\n in\n let started_at = Unix.gettimeofday () in\n let parent_id = Option.map Run.id parent in\n let manifest =\n make_manifest ~id ~experiment ~started_at ?name ?group ?parent ~tags ~params\n ~provenance ()\n in\n write_manifest (manifest_path dir) manifest;\n Index.add root ~id\n { experiment; name; group; parent_id; status = `running; tags; started_at };\n { root; experiment; id; dir; mutex = Mutex.create (); closed = false }\n\nlet finish ?(status = `finished) t () =\n with_lock t (fun () ->\n if not t.closed then (\n append_event t\n (Event_log.Finished\n {\n status = status_to_string status;\n ended_at = Unix.gettimeofday ();\n });\n Index.update_status t.root ~id:t.id (status :> Index.status);\n t.closed <- true))\n\nlet with_run ?root ~experiment ?name ?parent ?tags ?params ?notes ?capture_env f\n =\n let session =\n start ?root ~experiment ?name ?parent ?tags ?params ?notes ?capture_env ()\n in\n match f session with\n | value ->\n finish session ();\n value\n | exception exn ->\n finish ~status:`failed session ();\n raise exn\n\nlet resume run =\n if not (Run.resumable run) then invalid_arg err_not_resumable;\n let root = root_of_run_dir (Run.dir run) in\n let t =\n {\n root;\n experiment = Run.experiment_name run;\n id = Run.id run;\n dir = Run.dir run;\n mutex = Mutex.create ();\n closed = false;\n }\n in\n append_event t (Event_log.Resumed { at = Unix.gettimeofday () });\n Index.update_status root ~id:(Run.id run) `running;\n t\n\nlet run t =\n match Run.load ~root:t.root ~experiment:t.experiment ~id:t.id with\n | Some run -> run\n | None -> failwith err_missing_manifest\n\nlet set_notes t notes =\n with_lock t (fun () ->\n if not t.closed then append_event t (Event_log.Notes notes))\n\nlet log_metric t ~step ?timestamp key value =\n with_lock t (fun () ->\n if not t.closed then\n let timestamp =\n Option.value timestamp ~default:(Unix.gettimeofday ())\n in\n append_event t (Event_log.Metric { step; timestamp; key; value }))\n\nlet log_metrics t ~step ?timestamp pairs =\n with_lock t (fun () ->\n if not t.closed then\n let timestamp =\n Option.value timestamp ~default:(Unix.gettimeofday ())\n in\n List.iter\n (fun (key, value) ->\n append_event t (Event_log.Metric { step; timestamp; key; value }))\n pairs)\n\nlet define_metric t key ?(summary = `Last) ?step_metric ?goal () =\n with_lock t (fun () ->\n if not t.closed then\n append_event t\n (Event_log.Define_metric { key; summary; step_metric; goal }))\n\nlet rel_path_of ~run_dir abs_path =\n String.sub abs_path\n (String.length run_dir + 1)\n (String.length abs_path - String.length run_dir - 1)\n\nlet media_dest_path run_dir key step ext =\n let parts = String.split_on_char '/' key in\n let dir_parts, leaf =\n match List.rev parts with\n | [] -> ([], \"media\")\n | [ single ] -> ([], single)\n | last :: rest -> (List.rev rest, last)\n in\n let media_dir =\n List.fold_left Filename.concat (Filename.concat run_dir \"media\") dir_parts\n in\n let filename = Printf.sprintf \"%s_%d%s\" leaf step ext in\n (media_dir, Filename.concat media_dir filename)\n\nlet log_media t ~step ~key ~kind ~path =\n with_lock t (fun () ->\n if not t.closed then begin\n if not (Sys.file_exists path) then invalid_arg (err_media_missing ^ path);\n let ext = Filename.extension path in\n let media_dir, dest = media_dest_path t.dir key step ext in\n Fs.ensure_dir media_dir;\n Fs.copy_file path dest;\n let timestamp = Unix.gettimeofday () in\n append_event t\n (Event_log.Media\n {\n step;\n timestamp;\n key;\n kind;\n path = rel_path_of ~run_dir:t.dir dest;\n })\n end)\n\nlet log_table t ~step ~key ~columns ~rows =\n with_lock t (fun () ->\n if not t.closed then begin\n let json =\n Json_utils.json_obj\n [\n (\"columns\", Jsont.Json.list (List.map Jsont.Json.string columns));\n ( \"rows\",\n Jsont.Json.list\n (List.map\n (fun row -> Jsont.Json.list (List.map Value.to_json row))\n rows) );\n ]\n in\n let media_dir, dest = media_dest_path t.dir key step \".json\" in\n Fs.ensure_dir media_dir;\n Fs.write_file dest (Json_utils.json_to_string ~pretty:true json ^ \"\\n\");\n let timestamp = Unix.gettimeofday () in\n append_event t\n (Event_log.Media\n {\n step;\n timestamp;\n key;\n kind = `Table;\n path = rel_path_of ~run_dir:t.dir dest;\n })\n end)\n\nlet set_summary t values =\n with_lock t (fun () ->\n if not t.closed then append_event t (Event_log.Summary values))\n\nlet add_tags t tags =\n with_lock t (fun () ->\n if (not t.closed) && tags <> [] then append_event t (Event_log.Tags tags))\n\nlet log_artifact t ~name ~kind ~path ?(metadata = []) ?(aliases = []) () =\n with_lock t (fun () ->\n if t.closed then failwith err_closed_session;\n if not (Sys.file_exists path) then invalid_arg (err_path_missing ^ path);\n let digest = Fs.sha256_path path in\n let blob_rel_path =\n Filename.concat (Filename.concat \"blobs\" \"sha256\") digest\n in\n let blob_abs_path = Filename.concat t.root blob_rel_path in\n if not (Sys.file_exists blob_abs_path) then\n Fs.copy_tree path blob_abs_path;\n let payload : Artifact.payload =\n if Fs.is_directory path then `dir else `file\n in\n let json_metadata =\n List.map (fun (k, v) -> (k, Value.to_json v)) metadata\n in\n let artifact =\n Artifact.create ~root:t.root ~name ~kind ~payload ~digest\n ~path:blob_rel_path ~metadata:json_metadata ~aliases\n ~producer_run_id:(Some t.id)\n in\n append_event t\n (Event_log.Artifact_output { name; version = Artifact.version artifact });\n artifact)\n\nlet use_artifact t artifact =\n with_lock t (fun () ->\n if not t.closed then (\n Artifact.add_consumer ~root:t.root ~name:(Artifact.name artifact)\n ~version:(Artifact.version artifact)\n t.id;\n append_event t\n (Event_log.Artifact_input\n {\n name = Artifact.name artifact;\n version = Artifact.version artifact;\n })))\n"
},
{
"path": "packages/munin/lib/session.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Active run writers.\n\n A session is the append-only mutation boundary for a run. All writes go\n through the event log; no direct state mutation. *)\n\ntype t\n(** The type for active run sessions. *)\n\n(** {1:lifecycle Lifecycle} *)\n\nval start :\n ?root:string ->\n experiment:string ->\n ?name:string ->\n ?group:string ->\n ?parent:Run.t ->\n ?tags:string list ->\n ?params:(string * Value.t) list ->\n ?notes:string ->\n ?capture_env:string list ->\n ?command:string list ->\n ?cwd:string ->\n ?hostname:string ->\n ?pid:int ->\n ?git_commit:string ->\n ?git_dirty:bool ->\n ?env:(string * string) list ->\n unit ->\n t\n(** [start ~experiment ()] starts a new run session.\n\n Creates the store directory structure if it does not exist.\n\n - [root] defaults to [$RAVEN_TRACKING_DIR] or [$XDG_DATA_HOME/raven/munin].\n - [name] defaults to [None].\n - [tags] defaults to [[]].\n - [params] defaults to [[]].\n - [notes] defaults to [None].\n - [capture_env] is a list of environment variable names to capture into\n provenance. Defaults to [[]].\n - [command] defaults to [Sys.argv].\n - [cwd] defaults to [Sys.getcwd ()].\n - [hostname] defaults to [Unix.gethostname ()].\n - [pid] defaults to [Unix.getpid ()].\n - [git_commit] defaults to the HEAD commit detected from [cwd].\n - [git_dirty] defaults to the working tree status detected from [cwd].\n - [env] defaults to the variables captured via [capture_env]. *)\n\nval with_run :\n ?root:string ->\n experiment:string ->\n ?name:string ->\n ?parent:Run.t ->\n ?tags:string list ->\n ?params:(string * Value.t) list ->\n ?notes:string ->\n ?capture_env:string list ->\n (t -> 'a) ->\n 'a\n(** [with_run ~experiment f] starts a run, calls [f], and finishes the run as\n [`finished] on success or [`failed] on exception. The exception is re-raised\n after the run is closed.\n\n Optional arguments default as in {!start}. *)\n\nval resume : Run.t -> t\n(** [resume run] reopens an unfinished run for additional logging.\n\n Raises [Invalid_argument] if [Run.resumable run] is [false]. *)\n\nval run : t -> Run.t\n(** [run t] is the current materialized view of the run.\n\n Raises [Failure] if the run manifest is missing. *)\n\nval finish : ?status:[ `finished | `failed | `killed ] -> t -> unit -> unit\n(** [finish t ()] closes the run with the given final status.\n\n [status] defaults to [`finished]. The trailing [unit] argument allows\n partial application as a finalizer (e.g.\n [Fun.protect ~finally:(finish session)]). Calling [finish] on an\n already-closed session is a no-op. *)\n\n(** {1:scalars Scalars} *)\n\nval log_metric : t -> step:int -> ?timestamp:float -> string -> float -> unit\n(** [log_metric t ~step key value] appends a scalar metric sample.\n\n [timestamp] defaults to [Unix.gettimeofday ()]. Silently ignored if the\n session is closed. *)\n\nval log_metrics :\n t -> step:int -> ?timestamp:float -> (string * float) list -> unit\n(** [log_metrics t ~step pairs] appends multiple scalar metric samples\n atomically.\n\n [timestamp] defaults to [Unix.gettimeofday ()]. *)\n\n(** {1:metric_defs Metric definitions} *)\n\nval define_metric :\n t ->\n string ->\n ?summary:[ `Min | `Max | `Mean | `Last | `None ] ->\n ?step_metric:string ->\n ?goal:[ `Minimize | `Maximize ] ->\n unit ->\n unit\n(** [define_metric t key ()] declares how a metric should be summarised and\n plotted.\n\n - [summary] controls how the run summary value is computed from history:\n [`Min] (best for loss), [`Max] (best for accuracy), [`Mean], [`Last]\n (default), [`None] (no auto-summary).\n - [step_metric] specifies another metric as x-axis (e.g. [\"epoch\"]).\n Defaults to [None].\n - [goal] declares whether lower ([`Minimize]) or higher ([`Maximize]) is\n better, used by the TUI for \"best\" badges and by comparisons. Defaults to\n [None]. *)\n\n(** {1:media Media} *)\n\nval log_media :\n t ->\n step:int ->\n key:string ->\n kind:[ `Image | `Audio | `Table | `File ] ->\n path:string ->\n unit\n(** [log_media t ~step ~key ~kind ~path] copies [path] into the run's [media/]\n directory and appends a media event to the log.\n\n The file is stored at [/media/_.] where\n [] preserves the key's slash-delimited hierarchy as directories.\n [kind] is metadata for renderers; the TUI ignores media events. Silently\n ignored if the session is closed.\n\n @raise Invalid_argument if [path] does not exist. *)\n\nval log_table :\n t ->\n step:int ->\n key:string ->\n columns:string list ->\n rows:Value.t list list ->\n unit\n(** [log_table t ~step ~key ~columns ~rows] stores a table as JSON in the run's\n [media/] directory and appends a media event with [kind = `Table].\n\n The JSON file has the structure [{\"columns\": [...], \"rows\": [...]}]. Useful\n for confusion matrices, per-class metrics, data samples. *)\n\n(** {1:metadata Metadata} *)\n\nval set_notes : t -> string option -> unit\n(** [set_notes t note] replaces the run note. [None] clears it. *)\n\nval set_summary : t -> (string * Value.t) list -> unit\n(** [set_summary t values] merges summary values into the run. Later writes\n replace earlier values for the same key. *)\n\nval add_tags : t -> string list -> unit\n(** [add_tags t tags] appends tags to the run. Duplicate tags are ignored by\n readers. Empty lists are not written. *)\n\n(** {1:artifacts Artifacts} *)\n\nval log_artifact :\n t ->\n name:string ->\n kind:Artifact.kind ->\n path:string ->\n ?metadata:(string * Value.t) list ->\n ?aliases:string list ->\n unit ->\n Artifact.t\n(** [log_artifact t ~name ~kind ~path ()] stores [path] as a versioned artifact,\n records it as an output of [t], and returns the created version.\n\n - [metadata] defaults to [[]].\n - [aliases] defaults to [[]].\n\n Raises [Failure] if the session is closed. Raises [Invalid_argument] if\n [path] does not exist. *)\n\nval use_artifact : t -> Artifact.t -> unit\n(** [use_artifact t artifact] records [artifact] as an input of [t]. *)\n"
},
{
"path": "packages/munin/lib/store.ml",
"content": "type t = { root : string }\n\nlet root t = t.root\n\nlet open_ ?root () =\n let root = Option.value root ~default:(Env.root ()) in\n Fs.ensure_dir (Filename.concat root \"experiments\");\n Fs.ensure_dir (Filename.concat root \"artifacts\");\n Fs.ensure_dir (Filename.concat (Filename.concat root \"blobs\") \"sha256\");\n { root }\n\nlet list_experiments t = Fs.list_dirs (Filename.concat t.root \"experiments\")\n\n(* Index-based listing: filter on header fields, return lazy Run.t *)\nlet list_runs_indexed index ~root ?experiment ?status ?tag ?parent ?group () =\n Hashtbl.to_seq index |> List.of_seq\n |> List.filter_map (fun (id, (entry : Index.entry)) ->\n if\n Option.fold ~none:true ~some:(String.equal entry.experiment) experiment\n && Option.fold ~none:true ~some:(( = ) entry.status) status\n && Option.fold ~none:true\n ~some:(fun t -> List.exists (String.equal t) entry.tags)\n tag\n && Option.fold ~none:true\n ~some:(fun p -> entry.parent_id = Some p)\n parent\n && Option.fold ~none:true ~some:(fun g -> entry.group = Some g) group\n then Some (Run.load_from_index ~root id entry)\n else None)\n |> List.sort (fun a b -> String.compare (Run.id b) (Run.id a))\n\n(* Filesystem-based listing: fallback when no index *)\nlet list_runs_scan ~root ?experiment ?status ?tag ?parent ?group () =\n let runs =\n match experiment with\n | Some experiment ->\n Run.list ~root ~experiment ?status ?tag ?parent ?group ()\n | None ->\n Fs.list_dirs (Filename.concat root \"experiments\")\n |> List.concat_map (fun experiment ->\n Run.list ~root ~experiment ?status ?tag ?parent ?group ())\n in\n List.sort (fun a b -> String.compare (Run.id b) (Run.id a)) runs\n\nlet list_runs t ?experiment ?status ?tag ?parent ?group () =\n match Index.read t.root with\n | Some index ->\n list_runs_indexed index ~root:t.root ?experiment ?status ?tag ?parent\n ?group ()\n | None ->\n list_runs_scan ~root:t.root ?experiment ?status ?tag ?parent ?group ()\n\nlet find_run t id =\n match Index.read t.root with\n | Some index -> (\n match Hashtbl.find_opt index id with\n | Some entry -> Some (Run.load_from_index ~root:t.root id entry)\n | None -> None)\n | None ->\n List.find_map\n (fun experiment -> Run.load ~root:t.root ~experiment ~id)\n (list_experiments t)\n\nlet latest_run t ?experiment ?status ?tag ?group () =\n match list_runs t ?experiment ?status ?tag ?group () with\n | run :: _ -> Some run\n | [] -> None\n\nlet find_artifact t ~name ~version = Artifact.load ~root:t.root ~name ~version\n\nlet list_artifacts t ?name ?kind ?alias ?producer_run ?consumer_run () =\n Artifact.list ~root:t.root ?name ?kind ?alias ?producer_run ?consumer_run ()\n\nlet delete_run t run =\n Fs.remove_tree (Run.dir run);\n let exp_dir =\n Filename.concat\n (Filename.concat t.root \"experiments\")\n (Run.experiment_name run)\n in\n if Fs.list_dirs (Filename.concat exp_dir \"runs\") = [] then\n Fs.remove_tree exp_dir;\n Index.remove t.root ~id:(Run.id run)\n\nlet gc t =\n let blobs_dir = Filename.concat (Filename.concat t.root \"blobs\") \"sha256\" in\n let referenced = Hashtbl.create 64 in\n List.iter\n (fun artifact -> Hashtbl.replace referenced (Artifact.digest artifact) ())\n (list_artifacts t ());\n let removed = ref 0 in\n Fs.list_dirs blobs_dir\n |> List.iter (fun digest ->\n if not (Hashtbl.mem referenced digest) then (\n Fs.remove_tree (Filename.concat blobs_dir digest);\n incr removed));\n !removed\n"
},
{
"path": "packages/munin/lib/store.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Local tracking stores.\n\n A store is the local root directory containing experiments, runs, artifacts,\n and blobs. *)\n\ntype t\n(** The type for store handles. *)\n\nval open_ : ?root:string -> unit -> t\n(** [open_ ()] opens the local tracking store, creating its root directories if\n needed.\n\n [root] defaults to [$RAVEN_TRACKING_DIR] or [$XDG_DATA_HOME/raven/munin]. *)\n\nval root : t -> string\n(** [root t] is the absolute store root path. *)\n\n(** {1:runs Runs} *)\n\nval list_experiments : t -> string list\n(** [list_experiments t] is the experiment names stored in [t]. *)\n\nval list_runs :\n t ->\n ?experiment:string ->\n ?status:Run.status ->\n ?tag:string ->\n ?parent:string ->\n ?group:string ->\n unit ->\n Run.t list\n(** [list_runs t ()] is the persisted runs in [t], filtered when the optional\n selectors are provided. When [experiment] is omitted, searches across all\n experiments. Results are sorted by identifier descending (newest first). *)\n\nval find_run : t -> string -> Run.t option\n(** [find_run t id] is the run with identifier [id], if present. Searches across\n all experiments. *)\n\nval latest_run :\n t ->\n ?experiment:string ->\n ?status:Run.status ->\n ?tag:string ->\n ?group:string ->\n unit ->\n Run.t option\n(** [latest_run t ()] is the most recently started run matching the optional\n filters, by identifier ordering. *)\n\n(** {1:artifacts Artifacts} *)\n\nval find_artifact : t -> name:string -> version:string -> Artifact.t option\n(** [find_artifact t ~name ~version] is the named artifact version or alias, if\n present. *)\n\nval list_artifacts :\n t ->\n ?name:string ->\n ?kind:Artifact.kind ->\n ?alias:string ->\n ?producer_run:string ->\n ?consumer_run:string ->\n unit ->\n Artifact.t list\n(** [list_artifacts t ()] is the stored artifacts in [t], filtered when the\n optional selectors are provided. *)\n\n(** {1:maintenance Maintenance} *)\n\nval delete_run : t -> Run.t -> unit\n(** [delete_run t run] removes [run] and its event log from the store. Does not\n remove shared blobs. Removes the experiment directory if no runs remain. *)\n\nval gc : t -> int\n(** [gc t] removes unreferenced blobs. Returns the number removed. *)\n"
},
{
"path": "packages/munin/lib/sys/config/discover.ml",
"content": "module C = Configurator.V1\n\nlet () =\n C.main ~name:\"metrics-config\" (fun c ->\n let system = C.ocaml_config_var_exn c \"system\" in\n let c_flags, c_library_flags =\n if String.equal system \"macosx\" then\n (* macOS: link against IOKit and CoreFoundation frameworks *)\n ([], [ \"-framework\"; \"IOKit\"; \"-framework\"; \"CoreFoundation\" ])\n else\n (* Linux: no extra flags needed *)\n ([], [])\n in\n C.Flags.write_sexp \"c_flags.sexp\" c_flags;\n C.Flags.write_sexp \"c_library_flags.sexp\" c_library_flags)\n"
},
{
"path": "packages/munin/lib/sys/config/dune",
"content": "(executable\n (name discover)\n (libraries dune-configurator))\n"
},
{
"path": "packages/munin/lib/sys/dune",
"content": "(library\n (name munin_sys)\n (public_name munin.sys)\n (private_modules sysstat)\n (libraries unix str munin threads.posix)\n (foreign_stubs\n (language c)\n (flags\n (:standard\n (:include c_flags.sexp)))\n (names sysstat_stubs))\n (c_library_flags\n (:include c_library_flags.sexp)))\n\n(rule\n (targets c_flags.sexp c_library_flags.sexp)\n (action\n (run %{exe:config/discover.exe})))\n"
},
{
"path": "packages/munin/lib/sys/munin_sys.ml",
"content": "include Sysstat\n\ntype t = { stop : bool Atomic.t; thread : Thread.t }\n\nlet define session key = Session.define_metric session key ~summary:`Last ()\n\nlet start session ?(interval = 2.0) () =\n define session \"sys/cpu_user\";\n define session \"sys/cpu_system\";\n define session \"sys/mem_used_pct\";\n define session \"sys/mem_used_gb\";\n define session \"sys/proc_cpu_pct\";\n define session \"sys/proc_mem_mb\";\n define session \"sys/disk_read_mbs\";\n define session \"sys/disk_write_mbs\";\n define session \"sys/disk_util_pct\";\n let stop_flag = Atomic.make false in\n let prev_cpu = ref (Sysstat.Cpu.sample ()) in\n let prev_proc = ref (Sysstat.Proc.Self.sample ()) in\n let prev_disk = ref (Sysstat.Disk_io.sample ()) in\n let prev_time = ref (Unix.gettimeofday ()) in\n let step = ref 0 in\n let thread =\n Thread.create\n (fun () ->\n while not (Atomic.get stop_flag) do\n Thread.delay interval;\n if not (Atomic.get stop_flag) then begin\n incr step;\n let now = Unix.gettimeofday () in\n let dt = now -. !prev_time in\n (* System CPU *)\n let cpu = Sysstat.Cpu.sample () in\n let cpu_stats = Sysstat.Cpu.compute ~prev:!prev_cpu ~next:cpu in\n prev_cpu := cpu;\n (* System memory *)\n let mem = Sysstat.Mem.sample () in\n let mem_pct =\n Int64.to_float mem.used *. 100. /. Int64.to_float mem.total\n in\n let mem_gb = Int64.to_float mem.used /. 1_073_741_824. in\n (* Process stats *)\n let proc = Sysstat.Proc.Self.sample () in\n let proc_stats =\n Sysstat.Proc.Self.compute ~prev:!prev_proc ~next:proc ~dt\n ~num_cores:None\n in\n prev_proc := proc;\n (* Disk I/O *)\n let disk = Sysstat.Disk_io.sample () in\n let disk_stats =\n Sysstat.Disk_io.compute ~prev:!prev_disk ~next:disk ~dt\n in\n prev_disk := disk;\n prev_time := now;\n Session.log_metrics session ~step:!step\n [\n (\"sys/cpu_user\", cpu_stats.user);\n (\"sys/cpu_system\", cpu_stats.system);\n (\"sys/mem_used_pct\", mem_pct);\n (\"sys/mem_used_gb\", mem_gb);\n (\"sys/proc_cpu_pct\", proc_stats.cpu_percent);\n ( \"sys/proc_mem_mb\",\n Int64.to_float proc_stats.rss_bytes /. 1_048_576. );\n ( \"sys/disk_read_mbs\",\n disk_stats.read_bytes_per_sec /. 1_048_576. );\n ( \"sys/disk_write_mbs\",\n disk_stats.write_bytes_per_sec /. 1_048_576. );\n (\"sys/disk_util_pct\", disk_stats.utilization_percent);\n ]\n end\n done)\n ()\n in\n { stop = stop_flag; thread }\n\nlet stop t =\n if not (Atomic.get t.stop) then begin\n Atomic.set t.stop true;\n Thread.join t.thread\n end\n"
},
{
"path": "packages/munin/lib/sys/munin_sys.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** System monitoring.\n\n Provides stateless, poll-based system metrics sampling and background\n session monitoring. Each module samples instantaneous or cumulative values\n from the operating system. CPU, network, and disk I/O statistics are\n cumulative since boot and require two samples to compute usage percentages;\n memory statistics are instantaneous.\n\n {1:platform Platform support}\n\n Supported platforms: Linux and macOS. Platform-specific behavior is\n documented per module. Some metrics have limited availability on certain\n platforms (e.g., macOS CPU counters populate only user/nice/system/idle\n fields).\n\n {1:background Background monitoring}\n\n {!start} spawns a background thread that periodically samples CPU, memory,\n and process statistics and logs them as scalar metrics via {!Munin.Session}.\n\n Logged metrics (all with [sys/] prefix):\n\n {b System-wide:}\n - [sys/cpu_user] — user CPU percentage (0–100)\n - [sys/cpu_system] — system CPU percentage (0–100)\n - [sys/mem_used_pct] — memory usage percentage (0–100)\n - [sys/mem_used_gb] — memory used in GB\n\n {b Per-process:}\n - [sys/proc_cpu_pct] — process CPU percentage\n - [sys/proc_mem_mb] — process resident set size in MB\n\n {b Disk I/O:}\n - [sys/disk_read_mbs] — disk read rate in MB/s\n - [sys/disk_write_mbs] — disk write rate in MB/s\n - [sys/disk_util_pct] — disk utilization percentage *)\n\ninclude module type of Sysstat\n(** @inline *)\n\n(** {1 Background monitoring} *)\n\ntype t\n(** The type for background monitors. *)\n\nval start : Session.t -> ?interval:float -> unit -> t\n(** [start session ~interval ()] begins periodic system monitoring.\n\n All [sys/] metrics are defined with [~summary:`Last] so the final sampled\n value appears in run summaries.\n\n [interval] defaults to [2.0] seconds. The first sample is taken after one\n interval. The monitor thread is a daemon thread. *)\n\nval stop : t -> unit\n(** [stop t] signals the monitoring thread to exit and blocks until it\n terminates. Safe to call multiple times. *)\n"
},
{
"path": "packages/munin/lib/sys/sysstat.ml",
"content": "(* C binding types *)\n\ntype raw_memory = {\n total : int64;\n page_size : int64;\n active : int64;\n inactive : int64;\n speculative : int64;\n wired : int64;\n compressor : int64;\n purgeable : int64;\n external_ : int64;\n free : int64;\n swap_total : int64;\n swap_used : int64;\n}\n\ntype raw_network_io = {\n bytes_rx : int64;\n packets_rx : int64;\n bytes_tx : int64;\n packets_tx : int64;\n}\n\ntype raw_disk_io = {\n bytes_read : int64;\n bytes_written : int64;\n time_ms : int64;\n num_disks : int64;\n}\n\ntype raw_proc_info = {\n name : string;\n ppid : int;\n state : char;\n priority : int;\n nice : int;\n cmdline : string;\n user : string;\n user_time : int64;\n system_time : int64;\n resident_size : int64;\n virtual_size : int64;\n num_threads : int;\n num_running : int;\n faults : int64;\n}\n\ntype raw_statvfs = { total : int64; free : int64; avail : int64 }\ntype raw_self_mem = { rss : int64; vsize : int64 }\ntype raw_timebase = { numer : int; denom : int }\n\n(* C bindings *)\n\nexternal c_clk_tck : unit -> int = \"caml_sysstat_clk_tck\"\nexternal c_getpagesize : unit -> int = \"caml_sysstat_getpagesize\"\n\nexternal c_get_cpu_load : unit -> int64 array array\n = \"caml_sysstat_get_cpu_load\"\n\nexternal c_get_memory : unit -> raw_memory = \"caml_sysstat_get_memory\"\n\nexternal c_get_network_io : unit -> raw_network_io\n = \"caml_sysstat_get_network_io\"\n\nexternal c_get_disk_io : unit -> raw_disk_io = \"caml_sysstat_get_disk_io\"\nexternal c_list_pids : unit -> int array = \"caml_sysstat_list_pids\"\n\nexternal c_get_proc_info : int -> raw_proc_info option\n = \"caml_sysstat_get_proc_info\"\n\nexternal c_statvfs : string -> raw_statvfs = \"caml_sysstat_statvfs\"\nexternal c_proc_self_mem : unit -> raw_self_mem = \"caml_sysstat_proc_self_mem\"\nexternal c_get_timebase : unit -> raw_timebase = \"caml_sysstat_get_timebase\"\n\nexternal c_getmounts : unit -> (string * string * string) array\n = \"caml_sysstat_getmounts\"\n\nexternal c_get_loadavg : unit -> float * float * float\n = \"caml_sysstat_get_loadavg\"\n\nexternal c_get_uptime : unit -> int64 = \"caml_sysstat_get_uptime\"\n\n(* Helpers *)\n\nlet is_whitespace = function ' ' | '\\t' | '\\r' | '\\n' -> true | _ -> false\n\nlet split_whitespace (s : string) : string list =\n let len = String.length s in\n let rec skip i = if i < len && is_whitespace s.[i] then skip (i + 1) else i in\n let rec take i j =\n if j < len && not (is_whitespace s.[j]) then take i (j + 1) else j\n in\n let rec loop i acc =\n let i = skip i in\n if i >= len then List.rev acc\n else\n let j = take i i in\n loop j (String.sub s i (j - i) :: acc)\n in\n loop 0 []\n\nlet with_in_file path f =\n let ic = open_in path in\n match f ic with\n | v ->\n close_in_noerr ic;\n v\n | exception e ->\n close_in_noerr ic;\n raise e\n\nlet fold_lines path ~init ~f =\n with_in_file path (fun ic ->\n let rec loop acc =\n match input_line ic with\n | line -> loop (f acc line)\n | exception End_of_file -> acc\n in\n loop init)\n\nlet int64_of_string_opt s = try Some (Int64.of_string s) with _ -> None\nlet int_of_string_opt s = try Some (int_of_string s) with _ -> None\nlet i64_max0 x = if Int64.compare x 0L < 0 then 0L else x\n\nlet i64_delta ~prev ~next =\n let d = Int64.sub next prev in\n if Int64.compare d 0L < 0 then 0L else d\n\nlet f_delta_i64 ~prev ~next = Int64.to_float (i64_delta ~prev ~next)\nlet page_size_bytes : int64 Lazy.t = lazy (Int64.of_int (c_getpagesize ()))\nlet page_size () = Lazy.force page_size_bytes\n\n(* Time conversion *)\n\nlet clk_tck : int Lazy.t = lazy (c_clk_tck ())\n\nlet jiffies_to_ns (jiffies : int64) : int64 =\n let hz = Int64.of_int (Lazy.force clk_tck) in\n if Int64.compare hz 0L <= 0 then 0L\n else\n (* Avoid overflow: (q*1e9) + (r*1e9)/hz where jiffies = q*hz + r. *)\n let q = Int64.div jiffies hz in\n let r = Int64.rem jiffies hz in\n Int64.add\n (Int64.mul q 1_000_000_000L)\n (Int64.div (Int64.mul r 1_000_000_000L) hz)\n\nlet mach_timebase : (float * float) Lazy.t =\n lazy\n (let tb = c_get_timebase () in\n (float_of_int tb.numer, float_of_int tb.denom))\n\nlet mach_ticks_to_ns (ticks : int64) : float =\n let numer, denom = Lazy.force mach_timebase in\n Int64.to_float ticks *. numer /. denom\n\nlet proc_ticks_to_ns (ticks : int64) : float =\n if Sys.file_exists \"/proc\" then Int64.to_float (jiffies_to_ns ticks)\n else mach_ticks_to_ns ticks\n\n(* Cpu *)\n\nmodule Cpu = struct\n type t = {\n user : int64;\n nice : int64;\n system : int64;\n idle : int64;\n iowait : int64;\n irq : int64;\n softirq : int64;\n steal : int64;\n guest : int64;\n }\n\n type stats = {\n user : float;\n nice : float;\n system : float;\n idle : float;\n iowait : float;\n irq : float;\n softirq : float;\n steal : float;\n guest : float;\n }\n\n let zero : t =\n {\n user = 0L;\n nice = 0L;\n system = 0L;\n idle = 0L;\n iowait = 0L;\n irq = 0L;\n softirq = 0L;\n steal = 0L;\n guest = 0L;\n }\n\n let zero_stats =\n {\n user = 0.0;\n nice = 0.0;\n system = 0.0;\n idle = 100.0;\n iowait = 0.0;\n irq = 0.0;\n softirq = 0.0;\n steal = 0.0;\n guest = 0.0;\n }\n\n let of_linux_fields (fields : string list) : t option =\n match fields with\n | _name :: user :: nice :: system :: idle :: rest -> (\n match\n ( int64_of_string_opt user,\n int64_of_string_opt nice,\n int64_of_string_opt system,\n int64_of_string_opt idle )\n with\n | Some user, Some nice, Some system, Some idle ->\n let get idx =\n match List.nth_opt rest idx with\n | None -> 0L\n | Some s -> Option.value (int64_of_string_opt s) ~default:0L\n in\n Some\n {\n user;\n nice;\n system;\n idle;\n iowait = get 0;\n irq = get 1;\n softirq = get 2;\n steal = get 3;\n guest = get 4;\n }\n | _ -> None)\n | _ -> None\n\n let of_macos_row (row : int64 array) : t option =\n if Array.length row < 4 then None\n else\n Some\n {\n user = row.(0);\n nice = row.(1);\n system = row.(2);\n idle = row.(3);\n iowait = 0L;\n irq = 0L;\n softirq = 0L;\n steal = 0L;\n guest = 0L;\n }\n\n let is_cpu_core_name (name : string) =\n (* \"cpu0\", \"cpu1\", ... *)\n String.length name >= 4\n && String.starts_with ~prefix:\"cpu\" name\n && match name.[3] with '0' .. '9' -> true | _ -> false\n\n let read_linux () : (t * t list) option =\n let path = \"/proc/stat\" in\n if not (Sys.file_exists path) then None\n else\n let total = ref None in\n let cores = ref [] in\n let handle_line () line =\n match split_whitespace line with\n | [] -> ()\n | name :: _ as fields ->\n if name = \"cpu\" then total := of_linux_fields fields\n else if is_cpu_core_name name then\n match of_linux_fields fields with\n | Some t -> cores := t :: !cores\n | None -> ()\n else ()\n in\n (try\n fold_lines path ~init:() ~f:(fun () line ->\n handle_line () line;\n ())\n with _ -> ());\n match !total with None -> None | Some t -> Some (t, List.rev !cores)\n\n let read_macos () : (t * t list) option =\n try\n let rows = c_get_cpu_load () in\n if Array.length rows = 0 then None\n else\n match of_macos_row rows.(0) with\n | None -> None\n | Some total ->\n let cores =\n if Array.length rows <= 1 then []\n else\n Array.to_list\n (Array.init\n (Array.length rows - 1)\n (fun i ->\n match of_macos_row rows.(i + 1) with\n | Some t -> t\n | None -> zero))\n in\n Some (total, cores)\n with _ -> None\n\n let read () =\n if Sys.file_exists \"/proc/stat\" then read_linux () else read_macos ()\n\n let sample () : t =\n match read () with\n | Some (t, _) -> t\n | None -> raise (Sys_error \"Cpu.sample: failed to read CPU statistics\")\n\n let sample_per_core () : t array =\n match read () with\n | Some (_, cores) -> Array.of_list cores\n | None ->\n raise\n (Sys_error\n \"Cpu.sample_per_core: failed to read per-core CPU statistics\")\n\n let compute ~(prev : t) ~(next : t) : stats =\n let du = i64_delta ~prev:prev.user ~next:next.user in\n let dn = i64_delta ~prev:prev.nice ~next:next.nice in\n let ds = i64_delta ~prev:prev.system ~next:next.system in\n let di = i64_delta ~prev:prev.idle ~next:next.idle in\n let diw = i64_delta ~prev:prev.iowait ~next:next.iowait in\n let dirq = i64_delta ~prev:prev.irq ~next:next.irq in\n let dsi = i64_delta ~prev:prev.softirq ~next:next.softirq in\n let dst = i64_delta ~prev:prev.steal ~next:next.steal in\n let dg = i64_delta ~prev:prev.guest ~next:next.guest in\n\n let total =\n List.fold_left Int64.add 0L [ du; dn; ds; di; diw; dirq; dsi; dst; dg ]\n in\n if Int64.compare total 0L <= 0 then zero_stats\n else\n let total_f = Int64.to_float total in\n let pct v = Int64.to_float v /. total_f *. 100.0 in\n {\n user = pct du;\n nice = pct dn;\n system = pct ds;\n idle = pct di;\n iowait = pct diw;\n irq = pct dirq;\n softirq = pct dsi;\n steal = pct dst;\n guest = pct dg;\n }\nend\n\n(* Mem *)\n\nmodule Mem = struct\n type t = {\n total : int64;\n used : int64;\n free : int64;\n available : int64;\n compressed : int64;\n wired : int64;\n active : int64;\n inactive : int64;\n purgeable : int64;\n speculative : int64;\n external_ : int64;\n page_size : int64;\n swap_total : int64;\n swap_used : int64;\n }\n\n let read_linux_meminfo () =\n let path = \"/proc/meminfo\" in\n if not (Sys.file_exists path) then None\n else\n let mem_total = ref None in\n let mem_free = ref None in\n let mem_available = ref None in\n let buffers = ref None in\n let cached = ref None in\n let swap_total = ref None in\n let swap_free = ref None in\n let shmem = ref None in\n let sreclaimable = ref None in\n let set_kb (r : int64 option ref) (line : string) =\n match split_whitespace line with\n | _key :: value :: _ -> (\n match int64_of_string_opt value with\n | None -> ()\n | Some kb -> r := Some (Int64.mul kb 1024L))\n | _ -> ()\n in\n let handle_line line =\n if String.starts_with ~prefix:\"MemTotal:\" line then\n set_kb mem_total line\n else if String.starts_with ~prefix:\"MemFree:\" line then\n set_kb mem_free line\n else if String.starts_with ~prefix:\"MemAvailable:\" line then\n set_kb mem_available line\n else if String.starts_with ~prefix:\"Buffers:\" line then\n set_kb buffers line\n else if String.starts_with ~prefix:\"Cached:\" line then\n set_kb cached line\n else if String.starts_with ~prefix:\"SwapTotal:\" line then\n set_kb swap_total line\n else if String.starts_with ~prefix:\"SwapFree:\" line then\n set_kb swap_free line\n else if String.starts_with ~prefix:\"Shmem:\" line then set_kb shmem line\n else if String.starts_with ~prefix:\"SReclaimable:\" line then\n set_kb sreclaimable line\n else ()\n in\n (try fold_lines path ~init:() ~f:(fun () line -> handle_line line)\n with _ -> ());\n match !mem_total with\n | None -> None\n | Some total ->\n let free = Option.value !mem_free ~default:0L in\n let buffers_b = Option.value !buffers ~default:0L in\n let cached_b = Option.value !cached ~default:0L in\n let shmem_b = Option.value !shmem ~default:0L in\n let sreclaimable_b = Option.value !sreclaimable ~default:0L in\n let swap_total_b = Option.value !swap_total ~default:0L in\n let swap_free_b = Option.value !swap_free ~default:0L in\n\n let cache_used =\n i64_max0 (Int64.sub (Int64.add cached_b sreclaimable_b) shmem_b)\n in\n let used =\n i64_max0\n (Int64.sub\n (Int64.sub (Int64.sub total free) buffers_b)\n cache_used)\n in\n let swap_used = i64_max0 (Int64.sub swap_total_b swap_free_b) in\n (* MemAvailable is available on Linux 3.14+; fallback to free +\n buffers + cached *)\n let available =\n match !mem_available with\n | Some v -> v\n | None -> Int64.add free (Int64.add buffers_b cached_b)\n in\n\n Some\n {\n total;\n used;\n free;\n available;\n compressed = 0L;\n wired = 0L;\n active = 0L;\n inactive = cache_used;\n purgeable = buffers_b;\n speculative = 0L;\n external_ = 0L;\n page_size = page_size ();\n swap_total = swap_total_b;\n swap_used;\n }\n\n let read_macos () =\n try\n let raw = c_get_memory () in\n let ps = raw.page_size in\n let to_bytes pages = Int64.mul pages ps in\n let active = to_bytes raw.active in\n let inactive = to_bytes raw.inactive in\n let speculative = to_bytes raw.speculative in\n let wired = to_bytes raw.wired in\n let compressed = to_bytes raw.compressor in\n let purgeable = to_bytes raw.purgeable in\n let external_ = to_bytes raw.external_ in\n let free = to_bytes raw.free in\n\n let used_raw =\n Int64.add active\n (Int64.add inactive\n (Int64.add speculative (Int64.add wired compressed)))\n in\n let used_raw = i64_max0 (Int64.sub used_raw purgeable) in\n let used_raw = i64_max0 (Int64.sub used_raw external_) in\n let used_display = i64_max0 (Int64.sub used_raw compressed) in\n (* Approximate available as free + inactive + purgeable *)\n let available = Int64.add free (Int64.add inactive purgeable) in\n\n Some\n {\n total = raw.total;\n used = used_display;\n free;\n available;\n compressed;\n wired;\n active;\n inactive;\n purgeable;\n speculative;\n external_;\n page_size = ps;\n swap_total = raw.swap_total;\n swap_used = raw.swap_used;\n }\n with _ -> None\n\n let sample () : t =\n let result =\n if Sys.file_exists \"/proc/meminfo\" then read_linux_meminfo ()\n else read_macos ()\n in\n match result with\n | Some t -> t\n | None -> raise (Sys_error \"Mem.sample: failed to read memory statistics\")\nend\n\n(* Net *)\n\nmodule Net = struct\n type t = {\n bytes_rx : int64;\n packets_rx : int64;\n bytes_tx : int64;\n packets_tx : int64;\n }\n\n type stats = {\n rx_bytes_per_sec : float;\n rx_packets_per_sec : float;\n tx_bytes_per_sec : float;\n tx_packets_per_sec : float;\n }\n\n let read_linux () : t option =\n let path = \"/proc/net/dev\" in\n if not (Sys.file_exists path) then None\n else\n let add acc (rx_b, rx_p, tx_b, tx_p) =\n {\n bytes_rx = Int64.add acc.bytes_rx rx_b;\n packets_rx = Int64.add acc.packets_rx rx_p;\n bytes_tx = Int64.add acc.bytes_tx tx_b;\n packets_tx = Int64.add acc.packets_tx tx_p;\n }\n in\n let parse_line line =\n match String.split_on_char ':' line with\n | [ lhs; rhs ] -> (\n let iface = String.trim lhs in\n if iface = \"\" || iface = \"lo\" then None\n else\n let fields = split_whitespace rhs |> Array.of_list in\n if Array.length fields < 10 then None\n else\n match\n ( int64_of_string_opt fields.(0),\n int64_of_string_opt fields.(1),\n int64_of_string_opt fields.(8),\n int64_of_string_opt fields.(9) )\n with\n | Some rx_b, Some rx_p, Some tx_b, Some tx_p ->\n Some (rx_b, rx_p, tx_b, tx_p)\n | _ -> None)\n | _ -> None\n in\n let init =\n { bytes_rx = 0L; packets_rx = 0L; bytes_tx = 0L; packets_tx = 0L }\n in\n let total =\n fold_lines path ~init ~f:(fun acc line ->\n match parse_line line with\n | None -> acc\n | Some tuple -> add acc tuple)\n in\n Some total\n\n let read_macos () : t option =\n try\n let raw = c_get_network_io () in\n Some\n {\n bytes_rx = raw.bytes_rx;\n packets_rx = raw.packets_rx;\n bytes_tx = raw.bytes_tx;\n packets_tx = raw.packets_tx;\n }\n with _ -> None\n\n let sample () : t =\n let result =\n if Sys.file_exists \"/proc/net/dev\" then read_linux () else read_macos ()\n in\n match result with\n | Some t -> t\n | None -> raise (Sys_error \"Net.sample: failed to read network statistics\")\n\n let compute ~(prev : t) ~(next : t) ~(dt : float) : stats =\n if dt <= 0.0 then invalid_arg \"Net.compute: dt must be positive\";\n {\n rx_bytes_per_sec =\n f_delta_i64 ~prev:prev.bytes_rx ~next:next.bytes_rx /. dt;\n rx_packets_per_sec =\n f_delta_i64 ~prev:prev.packets_rx ~next:next.packets_rx /. dt;\n tx_bytes_per_sec =\n f_delta_i64 ~prev:prev.bytes_tx ~next:next.bytes_tx /. dt;\n tx_packets_per_sec =\n f_delta_i64 ~prev:prev.packets_tx ~next:next.packets_tx /. dt;\n }\nend\n\n(* Disk_io *)\n\nmodule Disk_io = struct\n type t = {\n bytes_read : int64;\n bytes_written : int64;\n time_ms : int64;\n num_disks : int64;\n }\n\n type stats = {\n read_bytes_per_sec : float;\n write_bytes_per_sec : float;\n utilization_percent : float;\n }\n\n type linux_disk = {\n name : string;\n sectors_read : int64;\n sectors_written : int64;\n io_time_ms : int64;\n }\n\n let is_excluded_name name =\n String.starts_with ~prefix:\"dm-\" name\n || String.starts_with ~prefix:\"loop\" name\n || String.starts_with ~prefix:\"md\" name\n || String.starts_with ~prefix:\"zram\" name\n\n let partition_base (name : string) : string option =\n (* If name ends with digits, it's *possibly* a partition. We return the\n candidate base device name. *)\n let len = String.length name in\n let rec find_first_non_digit i =\n if i < 0 then None\n else\n match name.[i] with\n | '0' .. '9' -> find_first_non_digit (i - 1)\n | _ -> Some i\n in\n match find_first_non_digit (len - 1) with\n | None -> None\n | Some i when i = len - 1 -> None (* no trailing digits *)\n | Some i ->\n let digit_start = i + 1 in\n if digit_start <= 0 then None\n else if name.[i] = 'p' then\n (* nvme0n1p1 / mmcblk0p1 style *)\n if i <= 0 then None else Some (String.sub name 0 i)\n else\n (* sda1 style *)\n Some (String.sub name 0 (i + 1))\n\n let read_linux () : t option =\n let path = \"/proc/diskstats\" in\n if not (Sys.file_exists path) then None\n else\n let parsed =\n fold_lines path ~init:[] ~f:(fun acc line ->\n match split_whitespace line with\n | _major :: _minor :: name :: _reads_ok :: _reads_merged\n :: sectors_read :: _time_read :: _writes_ok :: _writes_merged\n :: sectors_written :: _time_write :: _in_flight :: io_time_ms\n :: _weighted_time :: _ -> (\n if is_excluded_name name then acc\n else\n match\n ( int64_of_string_opt sectors_read,\n int64_of_string_opt sectors_written,\n int64_of_string_opt io_time_ms )\n with\n | Some sr, Some sw, Some t ->\n {\n name;\n sectors_read = sr;\n sectors_written = sw;\n io_time_ms = t;\n }\n :: acc\n | _ -> acc)\n | _ -> acc)\n in\n let candidates = List.rev parsed in\n let name_set = Hashtbl.create (List.length candidates) in\n List.iter (fun d -> Hashtbl.replace name_set d.name ()) candidates;\n\n let is_partition d =\n match partition_base d.name with\n | None -> false\n | Some base -> Hashtbl.mem name_set base\n in\n\n let bytes_read = ref 0L in\n let bytes_written = ref 0L in\n let time_ms = ref 0L in\n let num_disks = ref 0L in\n\n List.iter\n (fun d ->\n if not (is_partition d) then (\n bytes_read := Int64.add !bytes_read (Int64.mul d.sectors_read 512L);\n bytes_written :=\n Int64.add !bytes_written (Int64.mul d.sectors_written 512L);\n time_ms := Int64.add !time_ms d.io_time_ms;\n num_disks := Int64.add !num_disks 1L))\n candidates;\n\n Some\n {\n bytes_read = !bytes_read;\n bytes_written = !bytes_written;\n time_ms = !time_ms;\n num_disks = !num_disks;\n }\n\n let read_macos () : t option =\n try\n let raw = c_get_disk_io () in\n Some\n {\n bytes_read = raw.bytes_read;\n bytes_written = raw.bytes_written;\n time_ms = raw.time_ms;\n num_disks = raw.num_disks;\n }\n with _ -> None\n\n let sample () : t =\n let result =\n if Sys.file_exists \"/proc/diskstats\" then read_linux () else read_macos ()\n in\n match result with\n | Some t -> t\n | None ->\n raise (Sys_error \"Disk_io.sample: failed to read disk I/O statistics\")\n\n let compute ~(prev : t) ~(next : t) ~(dt : float) : stats =\n if dt <= 0.0 then invalid_arg \"Disk_io.compute: dt must be positive\";\n\n let read_bps =\n f_delta_i64 ~prev:prev.bytes_read ~next:next.bytes_read /. dt\n in\n let write_bps =\n f_delta_i64 ~prev:prev.bytes_written ~next:next.bytes_written /. dt\n in\n let time_delta_ms = f_delta_i64 ~prev:prev.time_ms ~next:next.time_ms in\n\n let utilization =\n if Int64.compare next.num_disks 0L <= 0 then 0.0\n else\n let denom = dt *. 1000.0 *. Int64.to_float next.num_disks in\n if denom <= 0.0 then 0.0 else 100.0 *. time_delta_ms /. denom\n in\n {\n read_bytes_per_sec = read_bps;\n write_bytes_per_sec = write_bps;\n utilization_percent = min utilization 100.0;\n }\nend\n\n(* Fs *)\n\nmodule Fs = struct\n type partition = {\n mount_point : string;\n total_bytes : int64;\n used_bytes : int64;\n avail_bytes : int64;\n }\n\n type t = {\n total_bytes : int64;\n used_bytes : int64;\n avail_bytes : int64;\n partitions : partition list;\n }\n\n let is_excluded_mount ~mount_point ~fstype =\n let excluded_fstypes =\n [\n \"devfs\";\n \"devtmpfs\";\n \"tmpfs\";\n \"proc\";\n \"sysfs\";\n \"cgroup\";\n \"autofs\";\n \"overlay\";\n ]\n in\n let excluded_mount_prefixes = [ \"/dev\"; \"/proc\"; \"/sys\"; \"/run\" ] in\n List.exists (fun p -> String.starts_with ~prefix:p fstype) excluded_fstypes\n || List.exists\n (fun p -> String.starts_with ~prefix:p mount_point)\n excluded_mount_prefixes\n\n let stat_partition (mount_point : string) : partition option =\n try\n let raw = c_statvfs mount_point in\n if Int64.compare raw.total 0L <= 0 then None\n else if Int64.compare raw.total 100_000_000L < 0 then None\n else\n let used = i64_max0 (Int64.sub raw.total raw.free) in\n Some\n {\n mount_point;\n total_bytes = raw.total;\n used_bytes = used;\n avail_bytes = raw.avail;\n }\n with _ -> None\n\n let partitions () : partition list =\n try\n c_getmounts () |> Array.to_list\n |> List.filter_map (fun (mount_point, _device, fstype) ->\n if is_excluded_mount ~mount_point ~fstype then None\n else stat_partition mount_point)\n with _ -> []\n\n let sample ?(path = \"/\") () : t =\n let result =\n try\n let raw = c_statvfs path in\n if Int64.compare raw.total 0L < 0 then None\n else\n let used = i64_max0 (Int64.sub raw.total raw.free) in\n Some\n {\n total_bytes = raw.total;\n used_bytes = used;\n avail_bytes = raw.avail;\n partitions = partitions ();\n }\n with _ -> None\n in\n match result with\n | Some t -> t\n | None ->\n raise\n (Sys_error\n (\"Fs.sample: failed to read filesystem statistics for \" ^ path))\nend\n\n(* Proc *)\n\nmodule Proc = struct\n type state =\n | Running\n | Sleeping\n | Disk_sleep\n | Stopped\n | Zombie\n | Idle\n | Unknown\n\n let state_of_char = function\n | 'R' -> Running\n | 'S' -> Sleeping\n | 'D' -> Disk_sleep\n | 'T' | 't' -> Stopped\n | 'Z' -> Zombie\n | 'I' -> Idle\n | 'X' -> Unknown (* dead *)\n | _ -> Unknown\n\n module Self = struct\n type t = {\n utime : float;\n stime : float;\n rss_bytes : int64;\n vsize_bytes : int64;\n }\n\n type stats = { cpu_percent : float; rss_bytes : int64; vsize_bytes : int64 }\n\n let read_proc_stat_tail (path : string) : string list option =\n try\n let line = with_in_file path input_line in\n match String.rindex_opt line ')' with\n | None -> None\n | Some k ->\n let start = k + 2 in\n if start >= String.length line then None\n else\n Some\n (split_whitespace\n (String.sub line start (String.length line - start)))\n with _ -> None\n\n let linux_rss_vsize () : (int64 * int64) option =\n match read_proc_stat_tail \"/proc/self/stat\" with\n | None -> None\n | Some fields -> (\n match (List.nth_opt fields 20, List.nth_opt fields 21) with\n | Some vsize_s, Some rss_pages_s -> (\n match\n (int64_of_string_opt vsize_s, int64_of_string_opt rss_pages_s)\n with\n | Some vsize, Some rss_pages ->\n let rss = Int64.mul rss_pages (page_size ()) in\n Some (rss, vsize)\n | _ -> None)\n | _ -> None)\n\n let sample () : t =\n let result =\n try\n let times = Unix.times () in\n let utime = times.Unix.tms_utime in\n let stime = times.Unix.tms_stime in\n\n let rss_bytes, vsize_bytes =\n if Sys.file_exists \"/proc/self/stat\" then\n match linux_rss_vsize () with\n | Some (rss, vsz) -> (rss, vsz)\n | None -> (0L, 0L)\n else\n let raw = c_proc_self_mem () in\n let rss = if Int64.compare raw.rss 0L < 0 then 0L else raw.rss in\n let vsz =\n if Int64.compare raw.vsize 0L < 0 then 0L else raw.vsize\n in\n (rss, vsz)\n in\n Some { utime; stime; rss_bytes; vsize_bytes }\n with _ -> None\n in\n match result with\n | Some t -> t\n | None ->\n raise\n (Sys_error \"Proc.Self.sample: failed to read process statistics\")\n\n let compute ~(prev : t) ~(next : t) ~(dt : float) ~(num_cores : int option)\n : stats =\n if dt <= 0.0 then invalid_arg \"Proc.Self.compute: dt must be positive\";\n\n let dt_ok = dt > 0.01 && dt < 10.0 in\n let cpu_delta = next.utime -. prev.utime +. (next.stime -. prev.stime) in\n let cpu_delta = max 0.0 cpu_delta in\n\n let cpu_percent =\n if not dt_ok then 0.0\n else\n let raw = cpu_delta /. dt *. 100.0 in\n match num_cores with\n | Some n when n > 0 && n <= 128 -> min (raw /. float_of_int n) 100.0\n | _ -> min raw 800.0\n in\n {\n cpu_percent;\n rss_bytes = next.rss_bytes;\n vsize_bytes = next.vsize_bytes;\n }\n end\n\n module Table = struct\n type t = {\n pid : int;\n ppid : int;\n name : string;\n cmdline : string;\n state : state;\n user : string;\n priority : int;\n nice : int;\n user_time : int64;\n system_time : int64;\n resident_size : int64;\n virtual_size : int64;\n num_threads : int;\n num_running : int;\n faults : int64;\n mem_percent : float;\n }\n\n type stats = {\n pid : int;\n name : string;\n cpu_percent : float;\n mem_percent : float;\n rss_bytes : int64;\n }\n\n let read_total_mem_bytes_linux () : int64 option =\n let path = \"/proc/meminfo\" in\n if not (Sys.file_exists path) then None\n else\n let mem_total = ref None in\n let handle line =\n if String.starts_with ~prefix:\"MemTotal:\" line then\n match split_whitespace line with\n | _key :: value :: _ -> (\n match int64_of_string_opt value with\n | Some kb -> mem_total := Some (Int64.mul kb 1024L)\n | None -> ())\n | _ -> ()\n in\n (try fold_lines path ~init:() ~f:(fun () line -> handle line)\n with _ -> ());\n !mem_total\n\n let total_mem_bytes () : int64 =\n if Sys.file_exists \"/proc/meminfo\" then\n Option.value (read_total_mem_bytes_linux ()) ~default:0L\n else try (c_get_memory ()).total with _ -> 0L\n\n let read_proc_stat_tail (path : string) : string list option =\n try\n let line = with_in_file path input_line in\n match String.rindex_opt line ')' with\n | None -> None\n | Some k ->\n let start = k + 2 in\n if start >= String.length line then None\n else\n Some\n (split_whitespace\n (String.sub line start (String.length line - start)))\n with _ -> None\n\n let read_comm (path : string) : string option =\n try\n let s = with_in_file path input_line |> String.trim in\n if s = \"\" then None else Some s\n with _ -> None\n\n let read_cmdline (pid : int) : string =\n let path = Printf.sprintf \"/proc/%d/cmdline\" pid in\n try\n let ic = open_in path in\n let len = in_channel_length ic in\n if len = 0 then (\n close_in_noerr ic;\n \"\")\n else\n let buf = Bytes.create len in\n let n = input ic buf 0 len in\n close_in_noerr ic;\n (* Replace null bytes with spaces *)\n for i = 0 to n - 1 do\n if Bytes.get buf i = '\\000' then Bytes.set buf i ' '\n done;\n String.trim (Bytes.sub_string buf 0 n)\n with _ -> \"\"\n\n let get_username (uid : int) : string =\n try (Unix.getpwuid uid).Unix.pw_name with _ -> \"\"\n\n let sample_linux () : t list =\n let total_mem = total_mem_bytes () in\n let dir = \"/proc\" in\n let dh = try Some (Unix.opendir dir) with _ -> None in\n match dh with\n | None -> []\n | Some dh ->\n let procs = ref [] in\n let add pid fields =\n (* Fields are tokens starting at /proc/[pid]/stat field 3 (state).\n Index 0 = state, 1 = ppid, ... 11 = utime, 12 = stime, 15 =\n priority, 16 = nice, 17 = num_threads, 20 = vsize, 21 = rss *)\n let state_char =\n match List.nth_opt fields 0 with\n | Some s when String.length s > 0 -> s.[0]\n | _ -> '?'\n in\n let ppid =\n match List.nth_opt fields 1 with\n | Some s -> Option.value (int_of_string_opt s) ~default:0\n | None -> 0\n in\n let priority =\n match List.nth_opt fields 15 with\n | Some s -> Option.value (int_of_string_opt s) ~default:0\n | None -> 0\n in\n let nice =\n match List.nth_opt fields 16 with\n | Some s -> Option.value (int_of_string_opt s) ~default:0\n | None -> 0\n in\n let num_threads =\n match List.nth_opt fields 17 with\n | Some s -> Option.value (int_of_string_opt s) ~default:1\n | None -> 1\n in\n match\n ( List.nth_opt fields 11,\n List.nth_opt fields 12,\n List.nth_opt fields 20,\n List.nth_opt fields 21 )\n with\n | Some ut_s, Some st_s, Some vsz_s, Some rss_pages_s -> (\n match\n ( int64_of_string_opt ut_s,\n int64_of_string_opt st_s,\n int64_of_string_opt vsz_s,\n int64_of_string_opt rss_pages_s )\n with\n | Some utime, Some stime, Some vsize, Some rss_pages ->\n let rss = Int64.mul rss_pages (page_size ()) in\n let mem_percent =\n if Int64.compare total_mem 0L > 0 then\n Int64.to_float rss *. 100.0 /. Int64.to_float total_mem\n else 0.0\n in\n let name =\n match read_comm (Printf.sprintf \"/proc/%d/comm\" pid) with\n | Some n -> n\n | None -> string_of_int pid\n in\n let cmdline = read_cmdline pid in\n let uid =\n try\n (Unix.stat (Printf.sprintf \"/proc/%d\" pid)).Unix.st_uid\n with _ -> -1\n in\n let user = if uid >= 0 then get_username uid else \"\" in\n procs :=\n {\n pid;\n ppid;\n name;\n cmdline;\n state = state_of_char state_char;\n user;\n priority;\n nice;\n user_time = utime;\n system_time = stime;\n resident_size = rss;\n virtual_size = vsize;\n num_threads;\n num_running = 0;\n faults = 0L;\n mem_percent;\n }\n :: !procs\n | _ -> ())\n | _ -> ()\n in\n (try\n while true do\n let entry = Unix.readdir dh in\n match int_of_string_opt entry with\n | None -> ()\n | Some pid -> (\n let stat_path = Printf.sprintf \"/proc/%d/stat\" pid in\n if Sys.file_exists stat_path then\n match read_proc_stat_tail stat_path with\n | Some fields -> add pid fields\n | None -> ())\n done\n with End_of_file -> ());\n Unix.closedir dh;\n !procs\n\n let sample_macos () : t list =\n let total_mem = total_mem_bytes () in\n let pids =\n try c_list_pids () |> Array.to_list |> List.filter (fun pid -> pid > 0)\n with _ -> []\n in\n let add acc pid =\n match c_get_proc_info pid with\n | None -> acc\n | Some info ->\n let mem_percent =\n if Int64.compare total_mem 0L > 0 then\n Int64.to_float info.resident_size\n *. 100.0 /. Int64.to_float total_mem\n else 0.0\n in\n {\n pid;\n ppid = info.ppid;\n name = info.name;\n cmdline = info.cmdline;\n state = state_of_char info.state;\n user = info.user;\n priority = info.priority;\n nice = info.nice;\n user_time = info.user_time;\n system_time = info.system_time;\n resident_size = info.resident_size;\n virtual_size = info.virtual_size;\n num_threads = info.num_threads;\n num_running = info.num_running;\n faults = info.faults;\n mem_percent;\n }\n :: acc\n in\n List.fold_left add [] pids\n\n let sample () : t list =\n try if Sys.file_exists \"/proc\" then sample_linux () else sample_macos ()\n with _ -> []\n\n let compute ~(prev : t list) ~(next : t list) ~(dt : float) : stats list =\n if dt <= 0.0 then invalid_arg \"Proc.Table.compute: dt must be positive\";\n\n let prev_by_pid : (int, t) Hashtbl.t =\n Hashtbl.create (List.length prev)\n in\n List.iter (fun (p : t) -> Hashtbl.replace prev_by_pid p.pid p) prev;\n\n let interval_ns = dt *. 1e9 in\n let cpu_percent (prev_p : t) (curr_p : t) =\n let prev_total = Int64.add prev_p.user_time prev_p.system_time in\n let curr_total = Int64.add curr_p.user_time curr_p.system_time in\n if Int64.compare curr_total prev_total <= 0 then 0.0\n else\n let delta_ticks = Int64.sub curr_total prev_total in\n let delta_ns = proc_ticks_to_ns delta_ticks in\n if interval_ns <= 0.0 then 0.0 else delta_ns /. interval_ns *. 100.0\n in\n\n let make_stats ~pid ~name ~cpu_percent ~mem_percent ~rss_bytes : stats =\n { pid; name; cpu_percent; mem_percent; rss_bytes }\n in\n next\n |> List.filter_map (fun (curr : t) ->\n match Hashtbl.find_opt prev_by_pid curr.pid with\n | Some prev_p ->\n let cpu = cpu_percent prev_p curr in\n if cpu > 0.0 || curr.mem_percent > 0.0 then\n Some\n (make_stats ~pid:curr.pid ~name:curr.name ~cpu_percent:cpu\n ~mem_percent:curr.mem_percent ~rss_bytes:curr.resident_size)\n else None\n | None ->\n if curr.mem_percent > 0.0 then\n Some\n (make_stats ~pid:curr.pid ~name:curr.name ~cpu_percent:0.0\n ~mem_percent:curr.mem_percent ~rss_bytes:curr.resident_size)\n else None)\n end\nend\n\n(* System info *)\n\nlet loadavg () = c_get_loadavg ()\nlet uptime () = c_get_uptime ()\n"
},
{
"path": "packages/munin/lib/sys/sysstat.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** System metrics collection library.\n\n Sysstat provides stateless, poll-based system monitoring. The caller manages\n state and sampling intervals. Each module samples instantaneous or\n cumulative values from the operating system. CPU, network, and disk I/O\n statistics are cumulative since boot and require two samples to compute\n usage percentages; memory statistics are instantaneous.\n\n {1:platform Platform support}\n\n Supported platforms: Linux and macOS. Platform-specific behavior is\n documented per module. Some metrics have limited availability on certain\n platforms (e.g., macOS CPU counters populate only user/nice/system/idle\n fields). *)\n\n(** {1:modules Modules} *)\n\n(** CPU statistics. *)\nmodule Cpu : sig\n type t = {\n user : int64; (** Time spent in user mode. *)\n nice : int64; (** Time spent in user mode with low priority (nice). *)\n system : int64; (** Time spent in system (kernel) mode. *)\n idle : int64; (** Time spent idle. *)\n iowait : int64; (** Time spent waiting for I/O to complete (Linux only). *)\n irq : int64; (** Time spent servicing interrupts (Linux only). *)\n softirq : int64; (** Time spent servicing soft interrupts (Linux only). *)\n steal : int64;\n (** Time stolen by other VMs in virtualized environments (Linux only).\n *)\n guest : int64; (** Time spent running guest VMs (Linux only). *)\n }\n (** Cumulative CPU time counters in ticks since boot.\n\n All fields represent cumulative time spent in each CPU state. The unit is\n platform-specific ticks (Linux jiffies or macOS Mach ticks) but is\n abstracted away by {!compute}.\n\n Platform behavior:\n - {b macOS}: Only [user], [nice], [system], and [idle] are available;\n other fields are [0L]. *)\n\n type stats = {\n user : float; (** Percentage of time spent in user mode. *)\n nice : float;\n (** Percentage of time spent in user mode with low priority. *)\n system : float; (** Percentage of time spent in system mode. *)\n idle : float; (** Percentage of time spent idle. *)\n iowait : float; (** Percentage of time spent waiting for I/O. *)\n irq : float; (** Percentage of time spent servicing interrupts. *)\n softirq : float; (** Percentage of time spent servicing soft interrupts. *)\n steal : float; (** Percentage of time stolen by hypervisor. *)\n guest : float; (** Percentage of time spent running guest VMs. *)\n }\n (** CPU usage percentages between two samples.\n\n All fields are in the range [0.0] to [100.0], where [100.0] represents\n full utilization. The sum of all fields equals [100.0]. *)\n\n val sample : unit -> t\n (** [sample ()] returns aggregate CPU counters across all cores.\n\n @raise Sys_error\n if CPU statistics are unavailable on the current platform or an error\n occurs during sampling. *)\n\n val sample_per_core : unit -> t array\n (** [sample_per_core ()] returns per-core CPU counters.\n\n The array length equals the number of logical CPU cores.\n\n @raise Sys_error\n if per-core statistics are unavailable or an error occurs. *)\n\n val compute : prev:t -> next:t -> stats\n (** [compute ~prev ~next] calculates CPU usage percentages between two\n samples.\n\n Computes the delta for each counter field and converts to percentages of\n total CPU time elapsed between samples. The sum of all fields in the\n returned [stats] equals [100.0].\n\n Returns all-zero percentages (with [idle = 100.0]) if no time has elapsed\n between samples (i.e., all counters are identical). This is not an error\n condition. *)\nend\n\n(** Memory statistics (instantaneous). *)\nmodule Mem : sig\n type t = {\n total : int64; (** Total physical memory. *)\n used : int64;\n (** Memory in use (calculated via platform-specific formula). *)\n free : int64; (** Free memory available for allocation. *)\n available : int64;\n (** Memory available for starting new applications without swapping.\n\n On macOS, approximated as [free + inactive + purgeable]. *)\n compressed : int64; (** Compressed memory (macOS only, [0L] on Linux). *)\n wired : int64;\n (** Wired (non-pageable) memory (macOS only, [0L] on Linux). *)\n active : int64; (** Active memory pages (macOS only, [0L] on Linux). *)\n inactive : int64;\n (** Inactive memory pages (Linux: cached + sreclaimable - shmem). *)\n purgeable : int64; (** Purgeable memory (macOS) or buffers (Linux). *)\n speculative : int64; (** Speculative memory (macOS only, [0L] on Linux). *)\n external_ : int64; (** External memory (macOS only, [0L] on Linux). *)\n page_size : int64; (** System page size in bytes. *)\n swap_total : int64; (** Total swap space. *)\n swap_used : int64; (** Used swap space. *)\n }\n (** Memory usage statistics.\n\n All size fields are in bytes. Memory statistics are instantaneous\n snapshots, not cumulative counters.\n\n Platform behavior:\n - {b Linux}: The [used] field is computed as\n [total - free - buffers - (cached + sreclaimable - shmem)].\n - {b macOS}: The [used] field is computed as\n [active + inactive + speculative + wired + compressed - purgeable -\n external], with [compressed] subtracted from the display value. *)\n\n val sample : unit -> t\n (** [sample ()] returns current memory usage statistics.\n\n @raise Sys_error\n if memory statistics are unavailable on the current platform or an error\n occurs during sampling. *)\nend\n\n(** Network I/O statistics. *)\nmodule Net : sig\n type t = {\n bytes_rx : int64; (** Total bytes received. *)\n packets_rx : int64; (** Total packets received. *)\n bytes_tx : int64; (** Total bytes transmitted. *)\n packets_tx : int64; (** Total packets transmitted. *)\n }\n (** Cumulative network I/O counters since boot.\n\n Aggregates all network interfaces except loopback. Counters are cumulative\n and monotonically increasing (until system reboot or counter overflow). *)\n\n type stats = {\n rx_bytes_per_sec : float; (** Receive rate in bytes per second. *)\n rx_packets_per_sec : float; (** Receive rate in packets per second. *)\n tx_bytes_per_sec : float; (** Transmit rate in bytes per second. *)\n tx_packets_per_sec : float; (** Transmit rate in packets per second. *)\n }\n (** Network I/O rates computed between two samples. *)\n\n val sample : unit -> t\n (** [sample ()] returns cumulative network I/O counters.\n\n @raise Sys_error\n if network statistics are unavailable on the current platform or an\n error occurs during sampling. *)\n\n val compute : prev:t -> next:t -> dt:float -> stats\n (** [compute ~prev ~next ~dt] calculates network I/O rates between two\n samples.\n\n Computes the delta for each counter and divides by [dt] to obtain rates\n per second. Negative deltas (e.g., from counter overflow or reboot) are\n treated as zero.\n\n @raise Invalid_argument if [dt <= 0.0]. *)\nend\n\n(** Disk I/O statistics. *)\nmodule Disk_io : sig\n type t = {\n bytes_read : int64; (** Total bytes read from disk. *)\n bytes_written : int64; (** Total bytes written to disk. *)\n time_ms : int64; (** Cumulative I/O time in milliseconds. *)\n num_disks : int64; (** Number of physical disks included in aggregation. *)\n }\n (** Cumulative disk I/O counters since boot.\n\n Aggregates physical disks only, excluding virtual devices, partitions, and\n metadata devices. Counters are cumulative and monotonically increasing\n until system reboot.\n\n Platform behavior:\n - {b Linux}: Excludes virtual devices ([dm-*], [loop*], [md*], [zram*])\n and partitions (detected by prefix matching the parent device name). *)\n\n type stats = {\n read_bytes_per_sec : float; (** Read rate in bytes per second. *)\n write_bytes_per_sec : float; (** Write rate in bytes per second. *)\n utilization_percent : float;\n (** Disk utilization percentage (0.0 to 100.0). *)\n }\n (** Disk I/O rates and utilization computed between two samples. *)\n\n val sample : unit -> t\n (** [sample ()] returns cumulative disk I/O counters.\n\n Only physical disks are included; partitions and virtual devices are\n excluded.\n\n @raise Sys_error\n if disk statistics are unavailable on the current platform or an error\n occurs during sampling. *)\n\n val compute : prev:t -> next:t -> dt:float -> stats\n (** [compute ~prev ~next ~dt] calculates disk I/O rates and utilization\n between two samples.\n\n Computes the delta for each counter and divides by [dt] to obtain rates\n per second. Utilization is calculated as\n [(time_delta / (dt * 1000 * num_disks)) * 100], representing the\n percentage of time disks were actively performing I/O, capped at [100.0].\n\n If [num_disks] is [0L], [utilization_percent] is [0.0].\n\n @raise Invalid_argument if [dt <= 0.0]. *)\nend\n\n(** Filesystem statistics (instantaneous). *)\nmodule Fs : sig\n type partition = {\n mount_point : string; (** Mount point path (e.g., [\"/\"], [\"/home\"]). *)\n total_bytes : int64; (** Total filesystem size. *)\n used_bytes : int64; (** Used space (calculated as [total - free]). *)\n avail_bytes : int64; (** Available space for unprivileged users. *)\n }\n (** Partition information for a single mounted filesystem.\n\n All size fields are in bytes. Represents a snapshot of filesystem usage at\n the time of sampling. *)\n\n type t = {\n total_bytes : int64; (** Total filesystem size for the queried path. *)\n used_bytes : int64; (** Used space for the queried path. *)\n avail_bytes : int64; (** Available space for the queried path. *)\n partitions : partition list;\n (** All mounted partitions (excluding virtual filesystems). *)\n }\n (** Filesystem statistics for a specific path.\n\n All size fields are in bytes. Contains statistics for the filesystem\n containing the specified path, plus a list of all mounted partitions. *)\n\n val sample : ?path:string -> unit -> t\n (** [sample ?path ()] returns filesystem statistics for the specified path.\n\n Returns statistics for the filesystem containing [path] (default: [\"/\"]),\n along with a list of all mounted partitions via {!partitions}.\n\n Virtual and system filesystems (e.g., [devfs], [tmpfs], [proc], [sysfs])\n are excluded from the partitions list. Filesystems smaller than\n approximately 100 MB are also excluded.\n\n @raise Sys_error if the path does not exist or an error occurs. *)\n\n val partitions : unit -> partition list\n (** [partitions ()] returns a list of all mounted partitions.\n\n Enumerates mounted filesystems and queries their usage. Excludes\n virtual/system filesystems and small filesystems (< 100 MB).\n\n Returns an empty list if no partitions are found or an error occurs. *)\nend\n\n(** Process statistics. *)\nmodule Proc : sig\n (** Process state. *)\n type state =\n | Running (** Currently executing on CPU. *)\n | Sleeping (** Interruptible sleep (waiting for event). *)\n | Disk_sleep (** Uninterruptible sleep (waiting for I/O). *)\n | Stopped (** Stopped by signal (e.g., SIGSTOP). *)\n | Zombie (** Terminated but not yet reaped by parent. *)\n | Idle (** Idle kernel thread. *)\n | Unknown (** State could not be determined. *)\n\n (** Current process (self) statistics. *)\n module Self : sig\n type t = {\n utime : float; (** Cumulative user-mode CPU time in seconds. *)\n stime : float; (** Cumulative system-mode CPU time in seconds. *)\n rss_bytes : int64; (** Resident set size (physical memory) in bytes. *)\n vsize_bytes : int64; (** Virtual memory size in bytes. *)\n }\n (** Raw process snapshot for delta calculation.\n\n Contains cumulative CPU time and instantaneous memory usage for the\n current process. CPU times are in seconds (converted from\n platform-specific units by [Unix.times]). *)\n\n type stats = {\n cpu_percent : float;\n (** CPU usage percentage (0.0 to 100.0 per core, or total if\n [num_cores] provided). *)\n rss_bytes : int64; (** Resident set size in bytes. *)\n vsize_bytes : int64; (** Virtual memory size in bytes. *)\n }\n (** Computed process statistics. *)\n\n val sample : unit -> t\n (** [sample ()] returns raw CPU times and memory usage for the current\n process.\n\n Uses [Unix.times] for CPU times.\n\n @raise Sys_error if an error occurs during sampling. *)\n\n val compute : prev:t -> next:t -> dt:float -> num_cores:int option -> stats\n (** [compute ~prev ~next ~dt ~num_cores] computes CPU usage percentage\n between two samples.\n\n CPU percentage is calculated as\n [((utime_delta + stime_delta) / dt) * 100]. If [num_cores] is provided,\n the percentage is normalized by dividing by the number of cores,\n yielding a value in [0.0] to [100.0]. Without normalization, the value\n can exceed [100.0] on multi-core systems.\n\n The result is clamped to prevent spurious values from timing anomalies:\n - With [num_cores]: capped at [100.0]\n - Without [num_cores]: capped at [800.0]\n\n If [dt] is outside the range [(0.01, 10.0)], returns [0.0] for\n [cpu_percent] to avoid division by near-zero or implausibly large\n intervals.\n\n @raise Invalid_argument if [dt <= 0.0]. *)\n end\n\n (** Process table statistics. *)\n module Table : sig\n type t = {\n pid : int; (** Process ID. *)\n ppid : int; (** Parent process ID. *)\n name : string; (** Process name (comm). *)\n cmdline : string;\n (** Full command line with arguments. Empty if unavailable. *)\n state : state; (** Current process state. *)\n user : string; (** Owner username. Empty if UID lookup fails. *)\n priority : int; (** Scheduling priority. *)\n nice : int; (** Nice value (-20 to 19). *)\n user_time : int64; (** Cumulative user-mode CPU time in ticks. *)\n system_time : int64; (** Cumulative system-mode CPU time in ticks. *)\n resident_size : int64; (** Resident set size in bytes. *)\n virtual_size : int64; (** Virtual memory size in bytes. *)\n num_threads : int; (** Number of threads. *)\n num_running : int;\n (** Number of running threads (macOS only, [0] on Linux). *)\n faults : int64; (** Page faults (macOS only, [0L] on Linux). *)\n mem_percent : float;\n (** Memory usage as percentage of total physical memory. *)\n }\n (** Raw process snapshot for delta calculation.\n\n Contains cumulative CPU time, state, and instantaneous memory/thread\n information for a process. CPU times are in platform-specific ticks\n (Linux jiffies or macOS Mach ticks). *)\n\n type stats = {\n pid : int; (** Process ID. *)\n name : string; (** Process name. *)\n cpu_percent : float; (** CPU usage percentage between samples. *)\n mem_percent : float;\n (** Memory usage as percentage of total physical memory. *)\n rss_bytes : int64; (** Resident set size in bytes. *)\n }\n (** Computed process statistics.\n\n Contains derived CPU percentage and filtered memory information. Only\n processes with non-zero CPU or memory usage are included. *)\n\n val sample : unit -> t list\n (** [sample ()] returns raw process snapshots for all running processes.\n\n Enumerates all processes visible to the current user and reads their\n statistics.\n\n Returns an empty list if an error occurs during enumeration. Individual\n process errors (e.g., process termination during sampling) are silently\n skipped. *)\n\n val compute : prev:t list -> next:t list -> dt:float -> stats list\n (** [compute ~prev ~next ~dt] calculates CPU percentages and filters\n processes.\n\n Matches processes by PID between [prev] and [next] samples. For matched\n processes, computes CPU percentage as:\n [(cpu_time_delta_ns / interval_ns) * 100].\n\n Only processes with non-zero [cpu_percent] or [mem_percent] are included\n in the result. New processes (in [next] but not [prev]) are included if\n their [mem_percent] is non-zero, with [cpu_percent] set to [0.0].\n\n @raise Invalid_argument if [dt <= 0.0]. *)\n end\nend\n\n(** {1:sysinfo System information} *)\n\nval loadavg : unit -> float * float * float\n(** [loadavg ()] returns the 1, 5, and 15 minute load averages. *)\n\nval uptime : unit -> int64\n(** [uptime ()] returns system uptime in seconds. *)\n"
},
{
"path": "packages/munin/lib/sys/sysstat_stubs.c",
"content": "/*\n * sysstat_stubs.c - OCaml FFI bindings for system metrics\n *\n * Platform-specific implementations.\n * macOS uses Mach APIs, IOKit, and libproc.\n * Linux uses /proc filesystem (parsed in OCaml), stubs return defaults.\n */\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#ifdef __APPLE__\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#endif\n\n#ifdef __linux__\n#include \n#endif\n\n/* Helpers */\n\n#ifdef __APPLE__\n/* Allocate a 4-element int64 array for CPU time counters */\nstatic value alloc_cpu_row(int64_t user, int64_t nice, int64_t sys,\n int64_t idle) {\n CAMLparam0();\n CAMLlocal1(arr);\n arr = caml_alloc(4, 0);\n Store_field(arr, 0, caml_copy_int64(user));\n Store_field(arr, 1, caml_copy_int64(nice));\n Store_field(arr, 2, caml_copy_int64(sys));\n Store_field(arr, 3, caml_copy_int64(idle));\n CAMLreturn(arr);\n}\n#endif\n\n/* CPU load */\n\nCAMLprim value caml_sysstat_get_cpu_load(value unit) {\n CAMLparam1(unit);\n CAMLlocal1(result);\n\n#ifdef __APPLE__\n CAMLlocal2(row, totalrow);\n\n natural_t ncpu = 0;\n processor_info_array_t cpuInfo;\n mach_msg_type_number_t numCpuInfo;\n\n host_t host = mach_host_self();\n kern_return_t kr = host_processor_info(host, PROCESSOR_CPU_LOAD_INFO, &ncpu,\n &cpuInfo, &numCpuInfo);\n mach_port_deallocate(mach_task_self(), host);\n\n if (kr != KERN_SUCCESS || ncpu == 0) {\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n\n int64_t sum_user = 0, sum_nice = 0, sum_sys = 0, sum_idle = 0;\n result = caml_alloc((mlsize_t)(ncpu + 1), 0);\n\n for (natural_t i = 0; i < ncpu; i++) {\n integer_t* base = cpuInfo + (CPU_STATE_MAX * i);\n int64_t user = (int64_t)base[CPU_STATE_USER];\n int64_t nice = (int64_t)base[CPU_STATE_NICE];\n int64_t sys = (int64_t)base[CPU_STATE_SYSTEM];\n int64_t idle = (int64_t)base[CPU_STATE_IDLE];\n\n sum_user += user;\n sum_nice += nice;\n sum_sys += sys;\n sum_idle += idle;\n\n row = alloc_cpu_row(user, nice, sys, idle);\n Store_field(result, (mlsize_t)(i + 1), row);\n }\n\n totalrow = alloc_cpu_row(sum_user, sum_nice, sum_sys, sum_idle);\n Store_field(result, 0, totalrow);\n\n vm_deallocate(mach_task_self(), (vm_address_t)cpuInfo,\n (vm_size_t)(numCpuInfo * sizeof(integer_t)));\n#else\n /* Linux: CPU stats read from /proc/stat in OCaml */\n result = caml_alloc(0, 0);\n#endif\n\n CAMLreturn(result);\n}\n\n/* Memory */\n\nCAMLprim value caml_sysstat_get_memory(value unit) {\n CAMLparam1(unit);\n CAMLlocal1(result);\n\n#ifdef __APPLE__\n host_t host = mach_host_self();\n\n /* Get total memory */\n host_basic_info_data_t hbi;\n mach_msg_type_number_t hbi_count = HOST_BASIC_INFO_COUNT;\n kern_return_t kr =\n host_info(host, HOST_BASIC_INFO, (host_info_t)&hbi, &hbi_count);\n if (kr != KERN_SUCCESS) {\n mach_port_deallocate(mach_task_self(), host);\n caml_failwith(\"host_info failed\");\n }\n\n /* Get VM statistics */\n vm_statistics64_data_t vm_stats;\n mach_msg_type_number_t vm_count = HOST_VM_INFO64_COUNT;\n kr = host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&vm_stats,\n &vm_count);\n mach_port_deallocate(mach_task_self(), host);\n if (kr != KERN_SUCCESS) {\n caml_failwith(\"host_statistics64 failed\");\n }\n\n /* Get swap usage */\n int mib[2] = {CTL_VM, VM_SWAPUSAGE};\n struct xsw_usage swap_info;\n size_t swap_size = sizeof(swap_info);\n int swap_ok = (sysctl(mib, 2, &swap_info, &swap_size, NULL, 0) == 0);\n\n /* Return raw page counts - OCaml computes derived values */\n result = caml_alloc_tuple(12);\n Store_field(result, 0, caml_copy_int64((int64_t)hbi.max_mem));\n Store_field(result, 1, caml_copy_int64((int64_t)vm_page_size));\n Store_field(result, 2, caml_copy_int64((int64_t)vm_stats.active_count));\n Store_field(result, 3, caml_copy_int64((int64_t)vm_stats.inactive_count));\n Store_field(result, 4, caml_copy_int64((int64_t)vm_stats.speculative_count));\n Store_field(result, 5, caml_copy_int64((int64_t)vm_stats.wire_count));\n Store_field(result, 6,\n caml_copy_int64((int64_t)vm_stats.compressor_page_count));\n Store_field(result, 7, caml_copy_int64((int64_t)vm_stats.purgeable_count));\n Store_field(result, 8,\n caml_copy_int64((int64_t)vm_stats.external_page_count));\n Store_field(result, 9, caml_copy_int64((int64_t)vm_stats.free_count));\n Store_field(result, 10,\n caml_copy_int64(swap_ok ? (int64_t)swap_info.xsu_total : 0));\n Store_field(result, 11,\n caml_copy_int64(swap_ok ? (int64_t)swap_info.xsu_used : 0));\n#else\n /* Linux: Memory stats read from /proc/meminfo in OCaml */\n result = caml_alloc_tuple(12);\n for (int i = 0; i < 12; i++) {\n Store_field(result, i, caml_copy_int64(0));\n }\n#endif\n\n CAMLreturn(result);\n}\n\n/* Network I/O */\n\nCAMLprim value caml_sysstat_get_network_io(value unit) {\n CAMLparam1(unit);\n CAMLlocal1(result);\n\n#ifdef __APPLE__\n int mib[6] = {CTL_NET, PF_ROUTE, 0, 0, NET_RT_IFLIST2, 0};\n size_t len = 0;\n char* buf = NULL;\n\n for (int retry = 0; retry < 4; retry++) {\n len = 0;\n if (sysctl(mib, 6, NULL, &len, NULL, 0) < 0 || len == 0) goto fail;\n\n len += 16 * retry * retry * sizeof(struct if_msghdr2);\n buf = malloc(len);\n if (!buf) goto fail;\n\n if (sysctl(mib, 6, buf, &len, NULL, 0) == 0) break;\n free(buf);\n buf = NULL;\n if (retry == 3) goto fail;\n }\n\n uint64_t bytes_rx = 0, packets_rx = 0, bytes_tx = 0, packets_tx = 0;\n for (char* next = buf; next < buf + len;) {\n struct if_msghdr* ifm = (struct if_msghdr*)next;\n next += ifm->ifm_msglen;\n if (ifm->ifm_type != RTM_IFINFO2) continue;\n\n struct if_msghdr2* ifm2 = (struct if_msghdr2*)ifm;\n if (ifm2->ifm_data.ifi_type != IFT_LOOP) {\n bytes_rx += ifm2->ifm_data.ifi_ibytes;\n packets_rx += ifm2->ifm_data.ifi_ipackets;\n bytes_tx += ifm2->ifm_data.ifi_obytes;\n packets_tx += ifm2->ifm_data.ifi_opackets;\n }\n }\n free(buf);\n\n result = caml_alloc_tuple(4);\n Store_field(result, 0, caml_copy_int64((int64_t)bytes_rx));\n Store_field(result, 1, caml_copy_int64((int64_t)packets_rx));\n Store_field(result, 2, caml_copy_int64((int64_t)bytes_tx));\n Store_field(result, 3, caml_copy_int64((int64_t)packets_tx));\n CAMLreturn(result);\n\nfail:\n result = caml_alloc_tuple(4);\n for (int i = 0; i < 4; i++) Store_field(result, i, caml_copy_int64(0));\n#else\n /* Linux: Network stats read from /proc/net/dev in OCaml */\n result = caml_alloc_tuple(4);\n for (int i = 0; i < 4; i++) Store_field(result, i, caml_copy_int64(0));\n#endif\n\n CAMLreturn(result);\n}\n\n/* Disk I/O */\n\nCAMLprim value caml_sysstat_get_disk_io(value unit) {\n CAMLparam1(unit);\n CAMLlocal1(result);\n\n#ifdef __APPLE__\n io_iterator_t drive_list;\n\n kern_return_t kr = IOServiceGetMatchingServices(\n kIOMainPortDefault, IOServiceMatching(\"IOBlockStorageDriver\"),\n &drive_list);\n\n if (kr != KERN_SUCCESS) {\n result = caml_alloc_tuple(4);\n for (int i = 0; i < 4; i++) Store_field(result, i, caml_copy_int64(0));\n CAMLreturn(result);\n }\n\n uint64_t read_sum = 0, write_sum = 0, time_sum = 0;\n uint64_t num_disks = 0;\n\n io_registry_entry_t drive;\n while ((drive = IOIteratorNext(drive_list)) != 0) {\n CFMutableDictionaryRef properties = NULL;\n\n if (IORegistryEntryCreateCFProperties(\n drive, &properties, kCFAllocatorDefault, 0) != KERN_SUCCESS) {\n IOObjectRelease(drive);\n continue;\n }\n\n if (!properties) {\n IOObjectRelease(drive);\n continue;\n }\n\n CFDictionaryRef statistics = CFDictionaryGetValue(\n properties, CFSTR(kIOBlockStorageDriverStatisticsKey));\n\n if (statistics) {\n num_disks++;\n CFNumberRef number;\n uint64_t value;\n\n number = CFDictionaryGetValue(\n statistics, CFSTR(kIOBlockStorageDriverStatisticsBytesReadKey));\n if (number) {\n CFNumberGetValue(number, kCFNumberSInt64Type, &value);\n read_sum += value;\n }\n\n number = CFDictionaryGetValue(\n statistics, CFSTR(kIOBlockStorageDriverStatisticsBytesWrittenKey));\n if (number) {\n CFNumberGetValue(number, kCFNumberSInt64Type, &value);\n write_sum += value;\n }\n\n number = CFDictionaryGetValue(\n statistics, CFSTR(kIOBlockStorageDriverStatisticsTotalReadTimeKey));\n if (number) {\n CFNumberGetValue(number, kCFNumberSInt64Type, &value);\n time_sum += value;\n }\n\n number = CFDictionaryGetValue(\n statistics, CFSTR(kIOBlockStorageDriverStatisticsTotalWriteTimeKey));\n if (number) {\n CFNumberGetValue(number, kCFNumberSInt64Type, &value);\n time_sum += value;\n }\n }\n\n CFRelease(properties);\n IOObjectRelease(drive);\n }\n\n IOObjectRelease(drive_list);\n\n result = caml_alloc_tuple(4);\n Store_field(result, 0, caml_copy_int64((int64_t)read_sum));\n Store_field(result, 1, caml_copy_int64((int64_t)write_sum));\n Store_field(result, 2,\n caml_copy_int64((int64_t)(time_sum / 1000000))); /* ns to ms */\n Store_field(result, 3, caml_copy_int64((int64_t)num_disks));\n#else\n /* Linux: Disk stats read from /proc/diskstats in OCaml */\n result = caml_alloc_tuple(4);\n for (int i = 0; i < 4; i++) {\n Store_field(result, i, caml_copy_int64(0));\n }\n#endif\n\n CAMLreturn(result);\n}\n\n/* Process list */\n\nCAMLprim value caml_sysstat_list_pids(value unit) {\n CAMLparam1(unit);\n CAMLlocal1(result);\n\n#ifdef __APPLE__\n int bufsize = proc_listpids(PROC_ALL_PIDS, 0, NULL, 0);\n if (bufsize <= 0) {\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n\n pid_t* pids = malloc(bufsize);\n if (!pids) {\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n\n int bytes_written = proc_listpids(PROC_ALL_PIDS, 0, pids, bufsize);\n if (bytes_written <= 0) {\n free(pids);\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n\n int num_pids = bytes_written / (int)sizeof(pid_t);\n\n result = caml_alloc(num_pids, 0);\n for (int i = 0; i < num_pids; i++) {\n Store_field(result, i, Val_int(pids[i]));\n }\n\n free(pids);\n#else\n /* Linux: PIDs enumerated from /proc in OCaml */\n result = caml_alloc(0, 0);\n#endif\n\n CAMLreturn(result);\n}\n\n/* Process info */\n\n#ifdef __APPLE__\n/* Map macOS process state to a character code matching Linux convention */\nstatic char macos_state_char(int p_stat) {\n switch (p_stat) {\n case SIDL:\n return 'I'; /* Idle (being created) */\n case SRUN:\n return 'R'; /* Running */\n case SSLEEP:\n return 'S'; /* Sleeping */\n case SSTOP:\n return 'T'; /* Stopped */\n case SZOMB:\n return 'Z'; /* Zombie */\n default:\n return '?';\n }\n}\n\n/* Get command line arguments for a process using sysctl KERN_PROCARGS2 */\nstatic int get_proc_cmdline(pid_t pid, char* buf, size_t bufsize) {\n if (bufsize == 0) return 0;\n buf[0] = '\\0';\n\n int mib[3] = {CTL_KERN, KERN_PROCARGS2, pid};\n size_t argmax = 0;\n\n /* Get maximum argument size */\n size_t argmax_size = sizeof(argmax);\n int argmax_mib[2] = {CTL_KERN, KERN_ARGMAX};\n if (sysctl(argmax_mib, 2, &argmax, &argmax_size, NULL, 0) != 0) {\n argmax = 65536; /* fallback */\n }\n\n char* procargs = malloc(argmax);\n if (!procargs) return 0;\n\n size_t size = argmax;\n if (sysctl(mib, 3, procargs, &size, NULL, 0) != 0) {\n free(procargs);\n return 0;\n }\n\n /* Skip argc (first 4 bytes) */\n if (size < sizeof(int)) {\n free(procargs);\n return 0;\n }\n\n int argc;\n memcpy(&argc, procargs, sizeof(int));\n char* p = procargs + sizeof(int);\n char* end = procargs + size;\n\n /* Skip executable path */\n while (p < end && *p != '\\0') p++;\n while (p < end && *p == '\\0') p++;\n\n /* Copy arguments up to bufsize, replacing nulls with spaces */\n size_t written = 0;\n int arg_count = 0;\n while (p < end && arg_count < argc && written < bufsize - 1) {\n if (*p == '\\0') {\n if (written > 0 && written < bufsize - 1) {\n buf[written++] = ' ';\n }\n arg_count++;\n p++;\n } else {\n buf[written++] = *p++;\n }\n }\n\n /* Trim trailing space if present */\n if (written > 0 && buf[written - 1] == ' ') {\n written--;\n }\n buf[written] = '\\0';\n\n free(procargs);\n return (int)written;\n}\n#endif\n\nCAMLprim value caml_sysstat_get_proc_info(value v_pid) {\n CAMLparam1(v_pid);\n\n#ifdef __APPLE__\n CAMLlocal5(result, name_str, cmdline_str, user_str, some);\n\n pid_t pid = Int_val(v_pid);\n\n /* Get task info for CPU times, memory, threads */\n struct proc_taskinfo pti;\n int ret = proc_pidinfo(pid, PROC_PIDTASKINFO, 0, &pti, PROC_PIDTASKINFO_SIZE);\n\n if (ret != PROC_PIDTASKINFO_SIZE) {\n CAMLreturn(Val_int(0)); /* None */\n }\n\n /* Get BSD info for ppid, state, nice, uid */\n struct proc_bsdinfo pbi;\n ret = proc_pidinfo(pid, PROC_PIDTBSDINFO, 0, &pbi, PROC_PIDTBSDINFO_SIZE);\n\n int ppid = 0;\n char state_char = '?';\n int priority = 0;\n int nice = 0;\n uid_t uid = 0;\n\n if (ret == PROC_PIDTBSDINFO_SIZE) {\n ppid = (int)pbi.pbi_ppid;\n state_char = macos_state_char(pbi.pbi_status);\n nice = (int)pbi.pbi_nice;\n uid = pbi.pbi_uid;\n }\n\n /* Get priority from kinfo_proc via sysctl */\n {\n struct kinfo_proc kp;\n size_t kp_size = sizeof(kp);\n int mib[4] = {CTL_KERN, KERN_PROC, KERN_PROC_PID, pid};\n if (sysctl(mib, 4, &kp, &kp_size, NULL, 0) == 0 && kp_size > 0) {\n priority = (int)kp.kp_proc.p_priority;\n }\n }\n\n /* Get process name from path */\n char name[MAXPATHLEN];\n ret = proc_pidpath(pid, name, sizeof(name));\n if (ret <= 0) {\n proc_name(pid, name, sizeof(name));\n } else {\n char* slash = strrchr(name, '/');\n if (slash) {\n memmove(name, slash + 1, strlen(slash + 1) + 1);\n }\n }\n\n /* Get command line */\n char cmdline[4096];\n get_proc_cmdline(pid, cmdline, sizeof(cmdline));\n\n /* Get username from UID */\n char username[256] = \"\";\n struct passwd* pw = getpwuid(uid);\n if (pw && pw->pw_name) {\n strncpy(username, pw->pw_name, sizeof(username) - 1);\n username[sizeof(username) - 1] = '\\0';\n }\n\n /* Build result tuple with 14 fields */\n result = caml_alloc_tuple(14);\n name_str = caml_copy_string(name);\n Store_field(result, 0, name_str); /* name */\n Store_field(result, 1, Val_int(ppid)); /* ppid */\n Store_field(result, 2, Val_int(state_char)); /* state (as char code) */\n Store_field(result, 3, Val_int(priority)); /* priority */\n Store_field(result, 4, Val_int(nice)); /* nice */\n cmdline_str = caml_copy_string(cmdline);\n Store_field(result, 5, cmdline_str); /* cmdline */\n user_str = caml_copy_string(username);\n Store_field(result, 6, user_str); /* user */\n Store_field(result, 7,\n caml_copy_int64((int64_t)pti.pti_total_user)); /* user_time */\n Store_field(result, 8,\n caml_copy_int64((int64_t)pti.pti_total_system)); /* system_time */\n Store_field(\n result, 9,\n caml_copy_int64((int64_t)pti.pti_resident_size)); /* resident_size */\n Store_field(\n result, 10,\n caml_copy_int64((int64_t)pti.pti_virtual_size)); /* virtual_size */\n Store_field(result, 11, Val_int(pti.pti_threadnum)); /* num_threads */\n Store_field(result, 12, Val_int(pti.pti_numrunning)); /* num_running */\n Store_field(result, 13,\n caml_copy_int64((int64_t)pti.pti_faults)); /* faults */\n\n some = caml_alloc(1, 0);\n Store_field(some, 0, result);\n CAMLreturn(some);\n#else\n /* Linux: Process info read from /proc/[pid] in OCaml */\n (void)v_pid;\n CAMLreturn(Val_int(0)); /* None */\n#endif\n}\n\n/* Mach timebase */\n\nCAMLprim value caml_sysstat_get_timebase(value unit) {\n CAMLparam1(unit);\n CAMLlocal1(result);\n\n#ifdef __APPLE__\n mach_timebase_info_data_t info;\n mach_timebase_info(&info);\n\n result = caml_alloc_tuple(2);\n Store_field(result, 0, Val_int(info.numer));\n Store_field(result, 1, Val_int(info.denom));\n#else\n /* Linux: Not needed, return 1:1 ratio */\n result = caml_alloc_tuple(2);\n Store_field(result, 0, Val_int(1));\n Store_field(result, 1, Val_int(1));\n#endif\n\n CAMLreturn(result);\n}\n\n/* Cross-platform */\n\n/* Get system load averages (1, 5, 15 minute) */\nCAMLprim value caml_sysstat_get_loadavg(value unit) {\n CAMLparam1(unit);\n CAMLlocal1(result);\n\n double loadavg[3];\n if (getloadavg(loadavg, 3) != 3) {\n loadavg[0] = loadavg[1] = loadavg[2] = 0.0;\n }\n\n result = caml_alloc_tuple(3);\n Store_field(result, 0, caml_copy_double(loadavg[0]));\n Store_field(result, 1, caml_copy_double(loadavg[1]));\n Store_field(result, 2, caml_copy_double(loadavg[2]));\n\n CAMLreturn(result);\n}\n\n/* Get system uptime in seconds */\nCAMLprim value caml_sysstat_get_uptime(value unit) {\n CAMLparam1(unit);\n\n#ifdef __APPLE__\n struct timeval boottime;\n size_t size = sizeof(boottime);\n int mib[2] = {CTL_KERN, KERN_BOOTTIME};\n\n if (sysctl(mib, 2, &boottime, &size, NULL, 0) != 0) {\n CAMLreturn(caml_copy_int64(0));\n }\n\n struct timeval now;\n gettimeofday(&now, NULL);\n\n int64_t uptime = (int64_t)(now.tv_sec - boottime.tv_sec);\n CAMLreturn(caml_copy_int64(uptime));\n#else\n /* Linux: read from /proc/uptime */\n FILE* f = fopen(\"/proc/uptime\", \"r\");\n if (!f) {\n CAMLreturn(caml_copy_int64(0));\n }\n\n double uptime_secs = 0.0;\n if (fscanf(f, \"%lf\", &uptime_secs) != 1) {\n uptime_secs = 0.0;\n }\n fclose(f);\n\n CAMLreturn(caml_copy_int64((int64_t)uptime_secs));\n#endif\n}\n\n/* Get system page size */\nCAMLprim value caml_sysstat_getpagesize(value unit) {\n CAMLparam1(unit);\n CAMLreturn(Val_long(sysconf(_SC_PAGESIZE)));\n}\n\n/* Get disk filesystem statistics using statvfs */\nCAMLprim value caml_sysstat_statvfs(value v_path) {\n CAMLparam1(v_path);\n CAMLlocal1(tup);\n\n const char* path = String_val(v_path);\n struct statvfs st;\n if (statvfs(path, &st) != 0) {\n caml_failwith(\"statvfs failed\");\n }\n\n uint64_t fr =\n (st.f_frsize != 0) ? (uint64_t)st.f_frsize : (uint64_t)st.f_bsize;\n uint64_t total = fr * (uint64_t)st.f_blocks;\n uint64_t free = fr * (uint64_t)st.f_bfree;\n uint64_t avail = fr * (uint64_t)st.f_bavail;\n\n tup = caml_alloc_tuple(3);\n Store_field(tup, 0, caml_copy_int64((int64_t)total));\n Store_field(tup, 1, caml_copy_int64((int64_t)free));\n Store_field(tup, 2, caml_copy_int64((int64_t)avail));\n CAMLreturn(tup);\n}\n\n/* Get process (self) memory statistics */\nCAMLprim value caml_sysstat_proc_self_mem(value unit) {\n CAMLparam1(unit);\n CAMLlocal1(tup);\n\n#ifdef __APPLE__\n mach_task_basic_info_data_t info;\n mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT;\n kern_return_t kr = task_info(mach_task_self(), MACH_TASK_BASIC_INFO,\n (task_info_t)&info, &count);\n if (kr != KERN_SUCCESS) {\n tup = caml_alloc_tuple(2);\n Store_field(tup, 0, caml_copy_int64(-1));\n Store_field(tup, 1, caml_copy_int64(-1));\n CAMLreturn(tup);\n }\n tup = caml_alloc_tuple(2);\n Store_field(tup, 0, caml_copy_int64((int64_t)info.resident_size));\n Store_field(tup, 1, caml_copy_int64((int64_t)info.virtual_size));\n#else\n /* Linux: Handled in OCaml via /proc/self/stat */\n tup = caml_alloc_tuple(2);\n Store_field(tup, 0, caml_copy_int64(-1));\n Store_field(tup, 1, caml_copy_int64(-1));\n#endif\n\n CAMLreturn(tup);\n}\n\n/* Get clock ticks per second (for Linux jiffies conversion) */\nCAMLprim value caml_sysstat_clk_tck(value unit) {\n CAMLparam1(unit);\n CAMLreturn(Val_long(sysconf(_SC_CLK_TCK)));\n}\n\n/* Mount enumeration */\n\n/* Get list of mounted filesystems.\n Returns array of (mount_point, device, fstype) tuples. */\nCAMLprim value caml_sysstat_getmounts(value unit) {\n CAMLparam1(unit);\n CAMLlocal3(result, tup, str);\n\n#ifdef __APPLE__\n struct statfs* mntbuf;\n int count = getmntinfo(&mntbuf, MNT_NOWAIT);\n if (count <= 0) {\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n\n result = caml_alloc(count, 0);\n for (int i = 0; i < count; i++) {\n tup = caml_alloc_tuple(3);\n str = caml_copy_string(mntbuf[i].f_mntonname);\n Store_field(tup, 0, str);\n str = caml_copy_string(mntbuf[i].f_mntfromname);\n Store_field(tup, 1, str);\n str = caml_copy_string(mntbuf[i].f_fstypename);\n Store_field(tup, 2, str);\n Store_field(result, i, tup);\n }\n#elif defined(__linux__)\n FILE* f = setmntent(\"/proc/mounts\", \"r\");\n if (!f) {\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n\n /* Read all entries into a temporary buffer to avoid TOCTOU race */\n int capacity = 64;\n int count = 0;\n char** dirs = malloc(capacity * sizeof(char*));\n char** devs = malloc(capacity * sizeof(char*));\n char** types = malloc(capacity * sizeof(char*));\n\n if (!dirs || !devs || !types) {\n free(dirs);\n free(devs);\n free(types);\n endmntent(f);\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n\n struct mntent* mnt;\n while ((mnt = getmntent(f)) != NULL) {\n if (count >= capacity) {\n capacity *= 2;\n char** new_dirs = realloc(dirs, capacity * sizeof(char*));\n char** new_devs = realloc(devs, capacity * sizeof(char*));\n char** new_types = realloc(types, capacity * sizeof(char*));\n if (!new_dirs || !new_devs || !new_types) {\n /* Clean up on allocation failure */\n for (int j = 0; j < count; j++) {\n free(dirs[j]);\n free(devs[j]);\n free(types[j]);\n }\n free(new_dirs ? new_dirs : dirs);\n free(new_devs ? new_devs : devs);\n free(new_types ? new_types : types);\n endmntent(f);\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n dirs = new_dirs;\n devs = new_devs;\n types = new_types;\n }\n dirs[count] = strdup(mnt->mnt_dir);\n devs[count] = strdup(mnt->mnt_fsname);\n types[count] = strdup(mnt->mnt_type);\n if (!dirs[count] || !devs[count] || !types[count]) {\n /* Clean up on strdup failure */\n for (int j = 0; j <= count; j++) {\n free(dirs[j]);\n free(devs[j]);\n free(types[j]);\n }\n free(dirs);\n free(devs);\n free(types);\n endmntent(f);\n result = caml_alloc(0, 0);\n CAMLreturn(result);\n }\n count++;\n }\n endmntent(f);\n\n /* Now allocate OCaml structures from the buffered data */\n result = caml_alloc(count, 0);\n for (int i = 0; i < count; i++) {\n tup = caml_alloc_tuple(3);\n str = caml_copy_string(dirs[i]);\n Store_field(tup, 0, str);\n str = caml_copy_string(devs[i]);\n Store_field(tup, 1, str);\n str = caml_copy_string(types[i]);\n Store_field(tup, 2, str);\n Store_field(result, i, tup);\n }\n\n /* Free temporary buffers */\n for (int i = 0; i < count; i++) {\n free(dirs[i]);\n free(devs[i]);\n free(types[i]);\n }\n free(dirs);\n free(devs);\n free(types);\n#else\n result = caml_alloc(0, 0);\n#endif\n\n CAMLreturn(result);\n}\n"
},
{
"path": "packages/munin/lib/tui/detail.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\nmodule Charts = Matrix_charts\n\n(* Exponential moving average: alpha in (0,1]. *)\nlet ema alpha history =\n match history with\n | [] -> []\n | (s0, v0) :: rest ->\n let rec loop acc prev = function\n | [] -> List.rev acc\n | (s, v) :: xs ->\n let smoothed = (alpha *. v) +. ((1. -. alpha) *. prev) in\n loop ((s, smoothed) :: acc) smoothed xs\n in\n (s0, v0) :: loop [] v0 rest\n\n(* Statistics *)\n\ntype stats = {\n last : float;\n best : float option;\n best_step : int option;\n mean : float;\n count : int;\n}\n\nlet compute_stats history ~best =\n match history with\n | [] -> None\n | _ ->\n let count = ref 0 in\n let sum = ref 0. in\n let last = ref 0. in\n List.iter\n (fun (_, v) ->\n incr count;\n sum := !sum +. v;\n last := v)\n history;\n let n = !count in\n let mean = !sum /. float_of_int n in\n let best_step =\n match best with\n | None -> None\n | Some bv ->\n let rec find = function\n | [] -> None\n | (s, v) :: _ when Float.equal v bv -> Some s\n | _ :: rest -> find rest\n in\n find history\n in\n Some { last = !last; best; best_step; mean; count = n }\n\nlet format_step step =\n if step < 1000 then string_of_int step\n else\n let s = string_of_int step in\n let len = String.length s in\n let buf = Buffer.create (len + ((len - 1) / 3)) in\n for i = 0 to len - 1 do\n if i > 0 && (len - i) mod 3 = 0 then Buffer.add_char buf ',';\n Buffer.add_char buf s.[i]\n done;\n Buffer.contents buf\n\nlet view_stats_row stats =\n let parts =\n [ Printf.sprintf \"Last: %.4g\" stats.last ]\n @ (match (stats.best, stats.best_step) with\n | Some bv, Some bs ->\n [ Printf.sprintf \"Best: %.4g (step %s)\" bv (format_step bs) ]\n | Some bv, None -> [ Printf.sprintf \"Best: %.4g\" bv ]\n | None, _ -> [])\n @ [\n Printf.sprintf \"Mean: %.4g\" stats.mean;\n Printf.sprintf \"Samples: %s\" (format_step stats.count);\n ]\n in\n box ~justify_content:Center ~align_items:Center\n ~size:{ width = pct 100; height = auto }\n [ text ~style:Theme.muted_style (String.concat \" \\u{00B7} \" parts) ]\n\nlet meta_style = Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:12) ()\n\nlet view_metric_def_row (metric_def : Munin.Run.metric_def) =\n let parts =\n (match metric_def.goal with\n | Some `Minimize -> [ \"Goal: minimize\" ]\n | Some `Maximize -> [ \"Goal: maximize\" ]\n | None -> [])\n @ (match metric_def.summary with\n | `Last -> []\n | `Min -> [ \"Summary: min\" ]\n | `Max -> [ \"Summary: max\" ]\n | `Mean -> [ \"Summary: mean\" ]\n | `None -> [ \"Summary: none\" ])\n @\n match metric_def.step_metric with\n | Some sm -> [ Printf.sprintf \"Step metric: %s\" sm ]\n | None -> []\n in\n if parts = [] then None\n else\n Some\n (box ~justify_content:Center ~align_items:Center\n ~size:{ width = pct 100; height = auto }\n [ text ~style:meta_style (String.concat \" \\u{00B7} \" parts) ])\n\n(* View *)\n\nlet view ~tag ~history_for_tag ~best ~size ~smooth ~metric_def =\n let history = history_for_tag tag in\n let display_history =\n match smooth with None -> history | Some alpha -> ema alpha history\n in\n let title =\n match Theme.last_value history with\n | None -> tag\n | Some v -> Printf.sprintf \"%s [%.4f]\" tag v\n in\n let title = if Option.is_some smooth then title ^ \" (EMA)\" else title in\n let stats = compute_stats history ~best in\n let stats_rows =\n (match stats with Some s -> [ view_stats_row s ] | None -> [])\n @\n match metric_def with\n | Some md -> (\n match view_metric_def_row md with Some row -> [ row ] | None -> [])\n | None -> []\n in\n box ~flex_direction:Column ~gap:(gap 1) ~align_items:Center ~size\n ([\n box ~border:true ~title ~padding:(padding 1)\n ~size:{ width = pct 100; height = pct 100 }\n ~flex_grow:1.0\n [\n canvas\n ~size:{ width = pct 100; height = pct 100 }\n (fun c ~delta:_ ->\n let width = Canvas.width c in\n let height = Canvas.height c in\n let grid = Canvas.grid c in\n match smooth with\n | None ->\n Theme.draw_metric_chart ~compact:false history grid ~width\n ~height\n | Some _ ->\n if history = [] then ()\n else\n let to_arr h =\n Array.of_list\n (List.map\n (fun (step, value) -> (float_of_int step, value))\n h)\n in\n let raw_data = to_arr history in\n let smooth_data = to_arr display_history in\n let chart =\n Charts.empty ()\n |> Charts.with_frame\n (Charts.manual_frame ~margins:(1, 1, 1, 4) ())\n |> Charts.with_axes\n ~x:\n (Charts.Axis.default |> Charts.Axis.with_ticks 6\n |> Charts.Axis.with_style Theme.axis_style)\n ~y:\n (Charts.Axis.default |> Charts.Axis.with_ticks 4\n |> Charts.Axis.with_style Theme.axis_style\n |> Charts.Axis.with_format (fun _ v ->\n Printf.sprintf \"%.4g\" v))\n |> Charts.with_grid\n (Charts.Gridlines.default\n |> Charts.Gridlines.with_style Theme.grid_style\n |> Charts.Gridlines.with_x true\n |> Charts.Gridlines.with_y true)\n |> Charts.line ~id:\"raw\" ~resolution:`Braille2x4\n ~style:\n (Ansi.Style.make\n ~fg:(Ansi.Color.grayscale ~level:8)\n ())\n ~x:fst ~y:snd raw_data\n |> Charts.line ~id:\"smooth\" ~resolution:`Braille2x4\n ~style:(Ansi.Style.make ~fg:Ansi.Color.cyan ())\n ~x:fst ~y:snd smooth_data\n in\n ignore (Charts.draw chart grid ~width ~height));\n ];\n ]\n @ stats_rows)\n"
},
{
"path": "packages/munin/lib/tui/dune",
"content": "(library\n (name munin_tui)\n (public_name munin.tui)\n (libraries\n munin\n unix\n mosaic\n mosaic.ui\n jsont\n jsont.bytesrw\n matrix.charts\n matrix.grid))\n"
},
{
"path": "packages/munin/lib/tui/footer.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\n\nlet hints = function\n | `Dashboard ->\n \"\\u{2190}\\u{2191}\\u{2193}\\u{2192} Select \\u{00B7} Enter Open \\u{00B7} <> \\\n Batch \\u{00B7} [] System \\u{00B7} i Info \\u{00B7} q Quit\"\n | `Detail smooth -> (\n match smooth with\n | Theme.Off -> \"S Smooth \\u{00B7} Esc back\"\n | Light | Medium | Heavy ->\n Printf.sprintf \"S Smooth [%d] \\u{00B7} Esc back\"\n (Theme.smooth_display smooth))\n | `Info -> \"q/Esc back\"\n\nlet view ~mode =\n box ~padding:(padding_xy 2 0) ~background:Theme.header_bg\n ~size:{ width = pct 100; height = auto }\n [ text ~style:Theme.muted_style (hints mode) ]\n"
},
{
"path": "packages/munin/lib/tui/header.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\n\n(* Styles *)\n\nlet label_style = Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:16) ()\nlet value_style = Ansi.Style.make ~bold:true ~fg:Ansi.Color.white ()\n\n(* Helpers *)\n\nlet format_elapsed secs =\n let secs = int_of_float secs in\n let h = secs / 3600 in\n let m = secs mod 3600 / 60 in\n let s = secs mod 60 in\n Printf.sprintf \"%02d:%02d:%02d\" h m s\n\nlet format_step step =\n if step < 1000 then string_of_int step\n else\n let s = string_of_int step in\n let len = String.length s in\n let buf = Buffer.create (len + ((len - 1) / 3)) in\n for i = 0 to len - 1 do\n if i > 0 && (len - i) mod 3 = 0 then Buffer.add_char buf ',';\n Buffer.add_char buf s.[i]\n done;\n Buffer.contents buf\n\n(* View *)\n\nlet tag_style =\n Ansi.Style.make\n ~fg:(Ansi.Color.grayscale ~level:18)\n ~bg:(Ansi.Color.grayscale ~level:5)\n ()\n\nlet view ~run_id ~run_name ~tags ~latest_epoch ~total_epochs ~latest_step\n ~elapsed_secs ~status =\n let color = Theme.status_color status in\n let epoch_items =\n match (latest_epoch, total_epochs) with\n | None, _ -> []\n | Some e, Some t ->\n [\n text ~style:label_style \"Epoch \";\n text ~style:value_style (Printf.sprintf \"%d/%d\" e t);\n ]\n | Some e, None ->\n [\n text ~style:label_style \"Epoch \";\n text ~style:value_style (string_of_int e);\n ]\n in\n let step_items =\n match latest_step with\n | None -> []\n | Some s ->\n [\n text ~style:label_style \"Step \";\n text ~style:value_style (format_step s);\n ]\n in\n let sep () = text ~style:Theme.muted_style \" \\u{00B7} \" in\n let stats =\n [ epoch_items; step_items ]\n |> List.filter (fun l -> l <> [])\n |> List.mapi (fun i items -> if i > 0 then sep () :: items else items)\n |> List.flatten\n in\n let stats =\n if stats <> [] then\n stats\n @ [ sep (); text ~style:Theme.muted_style (format_elapsed elapsed_secs) ]\n else [ text ~style:Theme.muted_style (format_elapsed elapsed_secs) ]\n in\n box ~padding:(padding_xy 2 0) ~flex_direction:Row ~gap:(gap 2)\n ~align_items:Center ~background:Theme.header_bg\n ~size:{ width = pct 100; height = auto }\n ([\n text ~style:value_style \"Munin\"; text ~style:Theme.muted_style \"\\u{2502}\";\n ]\n @ (match run_name with\n | Some name ->\n [\n text ~style:value_style name;\n text ~style:Theme.muted_style (Printf.sprintf \" (%s)\" run_id);\n ]\n | None -> [ text ~style:Theme.muted_style run_id ])\n @ List.map (fun t -> text ~style:tag_style (Printf.sprintf \" %s \" t)) tags\n @ stats\n @ [\n box ~flex_grow:1.0 ~size:{ width = auto; height = auto } [];\n text ~style:(Ansi.Style.make ~fg:color ()) \"\\u{25CF}\";\n text\n ~style:(Ansi.Style.make ~bold:true ~fg:color ())\n (Theme.status_label status);\n ])\n"
},
{
"path": "packages/munin/lib/tui/info.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\n\n(* Helpers *)\n\nlet format_timestamp t =\n let tm = Unix.localtime t in\n Printf.sprintf \"%04d-%02d-%02d %02d:%02d:%02d\" (1900 + tm.tm_year)\n (1 + tm.tm_mon) tm.tm_mday tm.tm_hour tm.tm_min tm.tm_sec\n\nlet format_elapsed secs =\n let secs = int_of_float secs in\n let h = secs / 3600 in\n let m = secs mod 3600 / 60 in\n let s = secs mod 60 in\n Printf.sprintf \"%02d:%02d:%02d\" h m s\n\nlet short_hash s = if String.length s > 7 then String.sub s 0 7 else s\n\n(* View *)\n\nlet view ~(run : Munin.Run.t) ~(status : Theme.run_status) ~elapsed_secs\n ~(metric_defs : (string * Munin.Run.metric_def) list)\n ~(latest_metrics : (string * Munin.Run.metric) list)\n ~(step_metrics : string list) ~(best_for_tag : string -> float option) =\n let prov = Munin.Run.provenance run in\n let status_color = Theme.status_color status in\n let run_section =\n [ Overview.section_header \"Run\" ]\n @ (match Munin.Run.name run with\n | Some name -> [ Overview.kv_row \"Name\" name ]\n | None -> [])\n @ [ Overview.kv_row \"ID\" (Munin.Run.id run) ]\n @ [ Overview.kv_row \"Experiment\" (Munin.Run.experiment_name run) ]\n @ (match Munin.Run.group run with\n | Some g -> [ Overview.kv_row \"Group\" g ]\n | None -> [])\n @ [\n box ~flex_direction:Row ~gap:(gap 1)\n ~size:{ width = pct 100; height = auto }\n [\n text ~style:Overview.key_style \"Status\";\n box ~flex_grow:1.0 ~size:{ width = auto; height = auto } [];\n text ~style:(Ansi.Style.make ~fg:status_color ()) \"\\u{25CF} \";\n text\n ~style:(Ansi.Style.make ~bold:true ~fg:status_color ())\n (Theme.status_label status);\n ];\n Overview.kv_row \"Started\" (format_timestamp (Munin.Run.started_at run));\n Overview.kv_row \"Duration\" (format_elapsed elapsed_secs);\n ]\n @\n match Munin.Run.tags run with\n | [] -> []\n | tags -> [ Overview.kv_row \"Tags\" (String.concat \", \" tags) ]\n in\n let prov_section =\n [ Overview.section_header \"Provenance\" ]\n @ [\n Overview.kv_row \"Command\" (String.concat \" \" prov.command);\n Overview.kv_row \"Directory\" prov.cwd;\n ]\n @ (match prov.hostname with\n | Some h -> [ Overview.kv_row \"Hostname\" h ]\n | None -> [])\n @ [ Overview.kv_row \"PID\" (string_of_int prov.pid) ]\n @\n match prov.git_commit with\n | Some hash ->\n let dirty =\n match prov.git_dirty with\n | Some true -> \" (dirty)\"\n | Some false -> \" (clean)\"\n | None -> \"\"\n in\n [ Overview.kv_row \"Git\" (short_hash hash ^ dirty) ]\n | None -> []\n in\n let params = Munin.Run.params run in\n let params_section =\n if params = [] then []\n else\n [ Overview.section_header \"Params\" ]\n @ List.map\n (fun (k, v) -> Overview.kv_row k (Overview.format_value v))\n params\n in\n let metrics =\n List.filter\n (fun (k, _) -> (not (Overview.is_sys k)) && not (List.mem k step_metrics))\n latest_metrics\n in\n let summary_section =\n if metrics = [] then []\n else\n [ Overview.section_header \"Summary\" ]\n @ List.map\n (fun (k, m) ->\n Overview.format_summary_value ~metric_defs ~best_for_tag k m)\n metrics\n in\n let notes_section =\n match Munin.Run.notes run with\n | Some notes when notes <> \"\" ->\n [\n Overview.section_header \"Notes\"; text ~style:Overview.val_style notes;\n ]\n | _ -> []\n in\n box ~flex_direction:Column ~padding:(padding 2) ~gap:(gap 1)\n ~size:{ width = pct 100; height = pct 100 }\n (run_section @ prov_section @ params_section @ summary_section\n @ notes_section)\n"
},
{
"path": "packages/munin/lib/tui/metrics.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\n\n(* Constants *)\n\nlet graph_height = 14\nlet min_graph_width = 25\nlet header_height = 3\nlet footer_height = 1\nlet metrics_padding = 2\n(* Layout helpers *)\n\nlet calculate_columns (available_width : int) : int =\n if available_width < min_graph_width * 2 then 1 else 2\n\nlet calculate_rows_per_batch (screen_height : int) : int =\n let available_height =\n screen_height - header_height - footer_height - (metrics_padding * 2)\n in\n if available_height < graph_height then 1\n else max 1 (available_height / graph_height)\n\nlet calculate_graphs_per_batch ~width ~height : int =\n let columns = calculate_columns width in\n let rows = calculate_rows_per_batch height in\n rows * columns\n\n(* Component state and update *)\n\ntype state = {\n screen_width : int;\n screen_height : int;\n current_batch : int;\n selected : int;\n}\n\ntype msg =\n | Resize of int * int\n | Next_batch\n | Prev_batch\n | Select_left\n | Select_right\n | Select_up\n | Select_down\n\ntype batch_window = {\n start_idx : int;\n end_idx : int;\n total_batches : int;\n current_batch : int;\n}\n\nlet batch_window ~width ~height ~current_batch ~total_metrics =\n if total_metrics = 0 then\n { start_idx = 0; end_idx = 0; total_batches = 1; current_batch = 0 }\n else\n let per_batch = calculate_graphs_per_batch ~width ~height in\n let total_batches = (total_metrics + per_batch - 1) / per_batch in\n let current_batch = min current_batch (max 0 (total_batches - 1)) in\n let start_idx = current_batch * per_batch in\n let end_idx = min (start_idx + per_batch) total_metrics in\n { start_idx; end_idx; total_batches; current_batch }\n\nlet initial_state () =\n { screen_width = 80; screen_height = 24; current_batch = 0; selected = 0 }\n\nlet visible_count s ~total_metrics =\n let w =\n batch_window ~width:s.screen_width ~height:s.screen_height\n ~current_batch:s.current_batch ~total_metrics\n in\n w.end_idx - w.start_idx\n\nlet update (msg : msg) (s : state) ~total_metrics : state =\n match msg with\n | Resize (width, height) ->\n let w =\n batch_window ~width ~height ~current_batch:s.current_batch\n ~total_metrics\n in\n let n = w.end_idx - w.start_idx in\n {\n screen_width = width;\n screen_height = height;\n current_batch = w.current_batch;\n selected = min s.selected (max 0 (n - 1));\n }\n | Next_batch ->\n let w =\n batch_window ~width:s.screen_width ~height:s.screen_height\n ~current_batch:(s.current_batch + 1) ~total_metrics\n in\n if w.current_batch = s.current_batch then s\n else { s with current_batch = w.current_batch; selected = 0 }\n | Prev_batch ->\n let prev = max 0 (s.current_batch - 1) in\n if prev = s.current_batch then s\n else { s with current_batch = prev; selected = 0 }\n | Select_left ->\n let n = visible_count s ~total_metrics in\n if n = 0 || s.selected = 0 then s\n else { s with selected = s.selected - 1 }\n | Select_right ->\n let n = visible_count s ~total_metrics in\n if n = 0 || s.selected >= n - 1 then s\n else { s with selected = s.selected + 1 }\n | Select_up ->\n let cols = calculate_columns s.screen_width in\n if s.selected >= cols then { s with selected = s.selected - cols } else s\n | Select_down ->\n let n = visible_count s ~total_metrics in\n let cols = calculate_columns s.screen_width in\n if s.selected + cols < n then { s with selected = s.selected + cols }\n else s\n\nlet selected_tag (s : state) ~total_metrics ~all_tags =\n let w =\n batch_window ~width:s.screen_width ~height:s.screen_height\n ~current_batch:s.current_batch ~total_metrics\n in\n let idx = w.start_idx + s.selected in\n if idx < w.end_idx then List.nth_opt all_tags idx else None\n\n(* Metric grouping *)\n\nlet group_prefix tag =\n match String.index_opt tag '/' with\n | Some i -> String.sub tag 0 i\n | None -> \"\"\n\nlet section_style =\n Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:12) ~dim:true ()\n\nlet view_group_header prefix =\n box\n ~size:{ width = pct 100; height = auto }\n [ text ~style:section_style (Printf.sprintf \"\\u{2500}\\u{2500} %s \" prefix) ]\n\nlet prefix_groups metrics =\n let rec go current_prefix acc group = function\n | [] ->\n let groups =\n if group = [] then acc else (current_prefix, List.rev group) :: acc\n in\n List.rev groups\n | ((_, tag) as item) :: rest ->\n let p = group_prefix tag in\n if String.equal p current_prefix then\n go current_prefix acc (item :: group) rest\n else\n let acc =\n if group = [] then acc else (current_prefix, List.rev group) :: acc\n in\n go p acc [ item ] rest\n in\n match metrics with\n | [] -> []\n | ((_, tag) as item) :: rest -> go (group_prefix tag) [] [ item ] rest\n\n(* Chart rendering *)\n\nlet dim_border = Ansi.Color.grayscale ~level:6\nlet selected_border = Ansi.Color.white\n\nlet view_metric_chart ~history_for_tag ~(best_for_tag : string -> float option)\n ~goal_for_tag ~columns ~selected tag =\n let history = history_for_tag tag in\n let best = best_for_tag tag in\n let width_pct = if columns = 1 then 100 else 49 in\n let goal_arrow =\n match goal_for_tag tag with\n | Some `Minimize -> \" \\u{2193}\"\n | Some `Maximize -> \" \\u{2191}\"\n | None -> \"\"\n in\n let title =\n match (Theme.last_value history, best) with\n | Some value, Some best_val ->\n Printf.sprintf \"%s [%.4f] best: %.4f%s\" tag value best_val goal_arrow\n | Some value, None -> Printf.sprintf \"%s [%.4f]\" tag value\n | None, _ -> tag\n in\n let border_color = if selected then selected_border else dim_border in\n box ~key:tag ~border:true ~border_color ~title ~padding:(padding 0)\n ~size:{ width = pct width_pct; height = px 14 }\n [\n canvas\n ~size:{ width = pct 100; height = pct 100 }\n (fun c ~delta:_ ->\n Theme.draw_metric_chart ~compact:true history (Canvas.grid c)\n ~width:(Canvas.width c) ~height:(Canvas.height c));\n ]\n\n(* View *)\n\nlet rec chunk_by n = function\n | [] -> []\n | lst ->\n let rec take k acc = function\n | [] -> (List.rev acc, [])\n | x :: xs ->\n if k = 0 then (List.rev acc, x :: xs)\n else take (k - 1) (x :: acc) xs\n in\n let group, rest = take n [] lst in\n group :: chunk_by n rest\n\nlet view (s : state) ~metric_tags ~history_for_tag ~best_for_tag ~goal_for_tag =\n if metric_tags = [] then\n box ~padding:(padding 1)\n ~size:{ width = pct 100; height = auto }\n [ text ~style:Theme.muted_style \" Waiting for metrics...\" ]\n else\n let columns = calculate_columns s.screen_width in\n let total_metrics = List.length metric_tags in\n let w =\n batch_window ~width:s.screen_width ~height:s.screen_height\n ~current_batch:s.current_batch ~total_metrics\n in\n let visible_metrics =\n List.mapi (fun i tag -> (i, tag)) metric_tags\n |> List.filter (fun (i, _) -> i >= w.start_idx && i < w.end_idx)\n |> List.mapi (fun local_idx (_, tag) -> (local_idx, tag))\n in\n let groups = prefix_groups visible_metrics in\n let batch_header =\n if w.total_batches > 1 then\n [\n box ~flex_direction:Row ~justify_content:Flex_end ~align_items:Center\n ~size:{ width = pct 100; height = auto }\n [\n text ~style:Theme.muted_style\n (Printf.sprintf \"Batch %d/%d\" (w.current_batch + 1)\n w.total_batches);\n ];\n ]\n else []\n in\n let charts =\n List.concat_map\n (fun (prefix, group_metrics) ->\n let header =\n if prefix <> \"\" then [ view_group_header prefix ] else []\n in\n let rows = chunk_by columns group_metrics in\n let chart_rows =\n List.mapi\n (fun row_idx row ->\n box\n ~key:(Printf.sprintf \"row-%s-%d\" prefix row_idx)\n ~flex_direction:Row ~gap:(gap 1)\n ~size:{ width = pct 100; height = auto }\n (List.map\n (fun (local_idx, tag) ->\n view_metric_chart ~history_for_tag ~best_for_tag\n ~goal_for_tag ~columns\n ~selected:(local_idx = s.selected) tag)\n row))\n rows\n in\n header @ chart_rows)\n groups\n in\n box ~flex_direction:Column ~padding:(padding_lrtb 1 1 1 0) ~gap:(gap 1)\n ~size:{ width = pct 100; height = auto }\n (batch_header @ charts)\n"
},
{
"path": "packages/munin/lib/tui/munin_tui.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\n\n(* Model *)\n\ntype model = {\n run : Munin.Run.t;\n monitor : Munin.Run_monitor.t;\n run_status : Theme.run_status;\n metrics_state : Metrics.state;\n mode : mode;\n smooth : Theme.smooth;\n show_system : bool;\n screen_width : int;\n was_live : bool;\n run_completed : bool;\n}\n\nand mode = Dashboard | Detail of string | Info\n\ntype msg =\n | Tick of float\n | Quit\n | Metrics_msg of Metrics.msg\n | Open_metric of string\n | Open_selected\n | Close_metric\n | Open_info\n | Close_info\n | Toggle_smooth\n | Toggle_system\n | Terminal_resize of int * int\n\n(* Helpers *)\n\nlet run_status_of_live_status :\n Munin.Run_monitor.live_status -> Theme.run_status = function\n | `Live -> Theme.Live\n | `Stopped -> Theme.Stopped\n | `Done `finished -> Theme.Done\n | `Done `failed -> Theme.Failed\n | `Done `killed -> Theme.Killed\n | `Done `running -> Theme.Live\n\nlet latest_step monitor =\n let ms = Munin.Run_monitor.metrics monitor in\n match ms with\n | [] -> None\n | _ ->\n Some\n (List.fold_left\n (fun acc (_, (m : Munin.Run.metric)) -> max acc m.step)\n 0 ms)\n\nlet latest_epoch monitor =\n let ms = Munin.Run_monitor.metrics monitor in\n match List.assoc_opt \"epoch\" ms with\n | Some (m : Munin.Run.metric) -> Some (int_of_float m.value)\n | None -> None\n\nlet total_epochs run =\n match Munin.Run.find_param run \"epochs\" with\n | Some v -> Munin.Value.to_int v\n | None -> None\n\nlet elapsed_secs m =\n let end_time =\n match m.run_status with\n | Theme.Live -> Unix.gettimeofday ()\n | Stopped | Done | Failed | Killed -> (\n match Munin.Run.ended_at m.run with\n | Some t -> t\n | None -> Unix.gettimeofday ())\n in\n end_time -. Munin.Run.started_at m.run\n\nlet metrics_width m =\n if m.show_system then int_of_float (float_of_int m.screen_width *. 0.66)\n else m.screen_width\n\nlet best_value monitor tag =\n Option.map\n (fun (b : Munin.Run.metric) -> b.value)\n (Munin.Run_monitor.best monitor tag)\n\n(* View *)\n\nlet divider () =\n box\n ~size:{ width = px 1; height = pct 100 }\n ~background:(Ansi.Color.grayscale ~level:8)\n [ text \" \" ]\n\nlet is_sys_metric tag = String.length tag > 4 && String.sub tag 0 4 = \"sys/\"\n\nlet step_metrics monitor =\n let defs = Munin.Run_monitor.metric_defs monitor in\n let from_defs =\n List.filter_map (fun (_, (d : Munin.Run.metric_def)) -> d.step_metric) defs\n in\n if List.mem \"epoch\" from_defs then from_defs else \"epoch\" :: from_defs\n\nlet user_metric_tags monitor =\n let sms = step_metrics monitor in\n List.filter_map\n (fun (tag, _) ->\n if is_sys_metric tag || List.mem tag sms then None else Some tag)\n (Munin.Run_monitor.metrics monitor)\n\nlet blocks =\n [|\n \"\\u{2581}\";\n \"\\u{2582}\";\n \"\\u{2583}\";\n \"\\u{2584}\";\n \"\\u{2585}\";\n \"\\u{2586}\";\n \"\\u{2587}\";\n \"\\u{2588}\";\n |]\n\nlet mini_sparkline history ~width =\n let values = List.map snd history in\n let n = List.length values in\n if n = 0 then \"\"\n else\n let arr =\n if n <= width then Array.of_list values\n else\n let a = Array.of_list values in\n Array.sub a (n - width) width\n in\n let len = Array.length arr in\n let lo = Array.fold_left min infinity arr in\n let hi = Array.fold_left max neg_infinity arr in\n let range = hi -. lo in\n let buf = Buffer.create (len * 3) in\n Array.iter\n (fun v ->\n let idx =\n if range = 0. then 3\n else int_of_float ((v -. lo) /. range *. 7.) |> max 0 |> min 7\n in\n Buffer.add_string buf blocks.(idx))\n arr;\n Buffer.contents buf\n\nlet spark_style = Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:10) ()\nlet bold_white = Ansi.Style.make ~bold:true ~fg:Ansi.Color.white ()\n\nlet view_summary_banner m =\n let elapsed = elapsed_secs m in\n let h = int_of_float elapsed / 3600 in\n let mi = int_of_float elapsed mod 3600 / 60 in\n let s = int_of_float elapsed mod 60 in\n let duration = Printf.sprintf \"%02d:%02d:%02d\" h mi s in\n let status_color = Theme.status_color m.run_status in\n let status_label = Theme.status_label m.run_status in\n let metric_tags = user_metric_tags m.monitor in\n let capped =\n if List.length metric_tags > 8 then\n List.filteri (fun i _ -> i < 8) metric_tags\n else metric_tags\n in\n let metric_entries =\n List.map\n (fun tag ->\n let history = Munin.Run_monitor.history m.monitor tag in\n let spark = mini_sparkline history ~width:8 in\n let value =\n match best_value m.monitor tag with\n | Some v -> Printf.sprintf \"%.4g\" v\n | None -> (\n match Theme.last_value history with\n | Some v -> Printf.sprintf \"%.4g\" v\n | None -> \"-\")\n in\n box ~flex_direction:Row ~gap:(gap 1)\n ~size:{ width = pct 50; height = auto }\n [\n text ~style:Theme.muted_style tag;\n text ~style:spark_style spark;\n text ~style:bold_white value;\n ])\n capped\n in\n let rec pairs = function\n | [] -> []\n | [ x ] -> [ [ x ] ]\n | x :: y :: rest -> [ x; y ] :: pairs rest\n in\n let metric_rows =\n List.map\n (fun row ->\n box ~flex_direction:Row ~size:{ width = pct 100; height = auto } row)\n (pairs metric_entries)\n in\n box ~border:true\n ~border_color:(Ansi.Color.grayscale ~level:8)\n ~title:(Printf.sprintf \" Run %s \" status_label)\n ~padding:(padding_xy 2 0)\n ~size:{ width = pct 100; height = auto }\n ([\n box ~flex_direction:Row ~gap:(gap 1)\n ~size:{ width = pct 100; height = auto }\n [\n text\n ~style:(Ansi.Style.make ~fg:status_color ())\n (Printf.sprintf \"%s in %s\" status_label duration);\n ];\n ]\n @ metric_rows)\n\nlet view_dashboard m =\n let all_metrics = Munin.Run_monitor.metrics m.monitor in\n let metric_tags = user_metric_tags m.monitor in\n let history_for_tag tag = Munin.Run_monitor.history m.monitor tag in\n let best_for_tag tag = best_value m.monitor tag in\n let goal_for_tag tag =\n match List.assoc_opt tag (Munin.Run_monitor.metric_defs m.monitor) with\n | Some (d : Munin.Run.metric_def) -> d.goal\n | None -> None\n in\n let metrics_pct = if m.show_system then 66 else 100 in\n let sys_latest tag =\n match List.assoc_opt tag all_metrics with\n | Some (m : Munin.Run.metric) -> m.value\n | None -> 0.0\n in\n let sys_values =\n System.\n {\n cpu_user = sys_latest \"sys/cpu_user\";\n cpu_system = sys_latest \"sys/cpu_system\";\n mem_pct = sys_latest \"sys/mem_used_pct\";\n mem_gb = sys_latest \"sys/mem_used_gb\";\n proc_cpu = sys_latest \"sys/proc_cpu_pct\";\n proc_mem_mb = sys_latest \"sys/proc_mem_mb\";\n disk_read_mbs = sys_latest \"sys/disk_read_mbs\";\n disk_write_mbs = sys_latest \"sys/disk_write_mbs\";\n disk_util_pct = sys_latest \"sys/disk_util_pct\";\n }\n in\n let right_panel =\n if m.show_system then\n [\n divider ();\n box ~flex_direction:Column ~padding:(padding_lrtb 1 1 1 0) ~gap:(gap 1)\n ~size:{ width = pct 34; height = auto }\n [\n System.view sys_values ~history_for_tag;\n Overview.view ~run:m.run ~latest_metrics:all_metrics\n ~step_metrics:(step_metrics m.monitor)\n ~metric_defs:(Munin.Run_monitor.metric_defs m.monitor)\n ~best_for_tag:(best_value m.monitor);\n ];\n ]\n else []\n in\n let banner = if m.run_completed then [ view_summary_banner m ] else [] in\n box ~flex_direction:Column\n ~size:{ width = pct 100; height = pct 100 }\n ([\n Header.view ~run_id:(Munin.Run.id m.run) ~run_name:(Munin.Run.name m.run)\n ~tags:(Munin.Run.tags m.run) ~latest_epoch:(latest_epoch m.monitor)\n ~total_epochs:(total_epochs m.run) ~latest_step:(latest_step m.monitor)\n ~elapsed_secs:(elapsed_secs m) ~status:m.run_status;\n ]\n @ banner\n @ [\n box ~flex_direction:Row ~flex_grow:1.0 ~flex_shrink:1.0\n ~overflow:{ x = Hidden; y = Hidden }\n ~size:{ width = pct 100; height = auto }\n ([\n box\n ~size:{ width = pct metrics_pct; height = auto }\n [\n Metrics.view m.metrics_state ~metric_tags ~history_for_tag\n ~best_for_tag ~goal_for_tag;\n ];\n ]\n @ right_panel);\n Footer.view ~mode:`Dashboard;\n ])\n\nlet view_detail m tag =\n let smooth_param = Theme.smooth_alpha m.smooth in\n let history_for_tag t = Munin.Run_monitor.history m.monitor t in\n let metric_def =\n List.assoc_opt tag (Munin.Run_monitor.metric_defs m.monitor)\n in\n box ~flex_direction:Column\n ~size:{ width = pct 100; height = pct 100 }\n [\n box ~flex_grow:1.0 ~justify_content:Center ~align_items:Center\n ~size:{ width = pct 100; height = pct 100 }\n [\n Detail.view ~tag ~history_for_tag ~best:(best_value m.monitor tag)\n ~smooth:smooth_param ~metric_def\n ~size:{ width = pct 80; height = pct 80 };\n ];\n Footer.view ~mode:(`Detail m.smooth);\n ]\n\nlet view_info m =\n let all_metrics = Munin.Run_monitor.metrics m.monitor in\n box ~flex_direction:Column\n ~size:{ width = pct 100; height = pct 100 }\n [\n Header.view ~run_id:(Munin.Run.id m.run) ~run_name:(Munin.Run.name m.run)\n ~tags:(Munin.Run.tags m.run) ~latest_epoch:(latest_epoch m.monitor)\n ~total_epochs:(total_epochs m.run) ~latest_step:(latest_step m.monitor)\n ~elapsed_secs:(elapsed_secs m) ~status:m.run_status;\n box ~flex_grow:1.0 ~flex_shrink:1.0 ~overflow:{ x = Hidden; y = Hidden }\n ~size:{ width = pct 100; height = auto }\n [\n Info.view ~run:m.run ~status:m.run_status\n ~elapsed_secs:(elapsed_secs m)\n ~metric_defs:(Munin.Run_monitor.metric_defs m.monitor)\n ~latest_metrics:all_metrics ~step_metrics:(step_metrics m.monitor)\n ~best_for_tag:(best_value m.monitor);\n ];\n Footer.view ~mode:`Info;\n ]\n\nlet view m =\n match m.mode with\n | Dashboard -> view_dashboard m\n | Detail tag -> view_detail m tag\n | Info -> view_info m\n\n(* TEA core *)\n\nlet init ~run =\n let monitor = Munin.Run_monitor.start run in\n Munin.Run_monitor.poll monitor;\n let run_status =\n run_status_of_live_status (Munin.Run_monitor.live_status monitor)\n in\n let metrics_state = Metrics.initial_state () in\n ( {\n run;\n monitor;\n run_status;\n metrics_state;\n mode = Dashboard;\n smooth = Theme.Off;\n show_system = true;\n screen_width = 80;\n was_live = run_status = Theme.Live;\n run_completed = false;\n },\n Cmd.none )\n\nlet update msg m =\n match msg with\n | Tick _dt ->\n Munin.Run_monitor.poll m.monitor;\n let run_status =\n run_status_of_live_status (Munin.Run_monitor.live_status m.monitor)\n in\n let run_completed =\n m.run_completed\n || m.was_live && run_status <> Theme.Live\n && run_status <> Theme.Stopped\n in\n ({ m with run_status; run_completed }, Cmd.none)\n | Terminal_resize (width, height) ->\n let m = { m with screen_width = width } in\n let mw = metrics_width m in\n let metrics_state' =\n Metrics.update\n (Metrics.Resize (mw, height))\n m.metrics_state\n ~total_metrics:(List.length (user_metric_tags m.monitor))\n in\n ({ m with metrics_state = metrics_state' }, Cmd.none)\n | Metrics_msg metrics_msg ->\n let total_metrics = List.length (user_metric_tags m.monitor) in\n let metrics_state' =\n Metrics.update metrics_msg m.metrics_state ~total_metrics\n in\n ({ m with metrics_state = metrics_state' }, Cmd.none)\n | Open_metric tag -> ({ m with mode = Detail tag }, Cmd.none)\n | Open_selected -> (\n let all_tags = user_metric_tags m.monitor in\n let total_metrics = List.length all_tags in\n match Metrics.selected_tag m.metrics_state ~total_metrics ~all_tags with\n | Some tag -> ({ m with mode = Detail tag }, Cmd.none)\n | None -> (m, Cmd.none))\n | Close_metric -> ({ m with mode = Dashboard }, Cmd.none)\n | Open_info -> ({ m with mode = Info }, Cmd.none)\n | Close_info -> ({ m with mode = Dashboard }, Cmd.none)\n | Toggle_smooth -> (\n match m.mode with\n | Detail _ -> ({ m with smooth = Theme.next_smooth m.smooth }, Cmd.none)\n | Dashboard | Info -> (m, Cmd.none))\n | Toggle_system ->\n let m = { m with show_system = not m.show_system } in\n let mw = metrics_width m in\n let total_metrics = List.length (user_metric_tags m.monitor) in\n let metrics_state' =\n Metrics.update\n (Metrics.Resize (mw, m.metrics_state.screen_height))\n m.metrics_state ~total_metrics\n in\n ({ m with metrics_state = metrics_state' }, Cmd.none)\n | Quit ->\n Munin.Run_monitor.close m.monitor;\n (m, Cmd.quit)\n\nlet subscriptions m =\n Sub.batch\n [\n Sub.on_tick (fun ~dt -> Tick dt);\n Sub.on_resize (fun ~width ~height -> Terminal_resize (width, height));\n Sub.on_key (fun ev ->\n let is c ch =\n let lo = Uchar.of_char ch in\n let hi = Uchar.of_char (Char.uppercase_ascii ch) in\n Uchar.equal c lo || Uchar.equal c hi\n in\n let data = Mosaic_ui.Event.Key.data ev in\n match (data.key, m.mode) with\n | Char c, Detail _ when is c 's' -> Some Toggle_smooth\n | Char c, Dashboard when is c 'q' -> Some Quit\n | Char c, Detail _ when is c 'q' -> Some Close_metric\n | Char c, Info when is c 'q' -> Some Close_info\n | Char c, Dashboard when is c 'i' -> Some Open_info\n | Escape, Dashboard -> Some Quit\n | Escape, Detail _ -> Some Close_metric\n | Escape, Info -> Some Close_info\n | Left, Dashboard -> Some (Metrics_msg Metrics.Select_left)\n | Right, Dashboard -> Some (Metrics_msg Metrics.Select_right)\n | Up, Dashboard -> Some (Metrics_msg Metrics.Select_up)\n | Down, Dashboard -> Some (Metrics_msg Metrics.Select_down)\n | Char c, Dashboard when is c '[' -> Some Toggle_system\n | Char c, Dashboard when is c ']' -> Some Toggle_system\n | Char c, Dashboard when Uchar.equal c (Uchar.of_char '<') ->\n Some (Metrics_msg Metrics.Prev_batch)\n | Char c, Dashboard when Uchar.equal c (Uchar.of_char '>') ->\n Some (Metrics_msg Metrics.Next_batch)\n | Enter, Dashboard -> Some Open_selected\n | Char c, Dashboard when is c ' ' -> Some Open_selected\n | _ -> None);\n ]\n\nlet run ?root ?experiment ?runs () =\n let store = Munin.Store.open_ ?root () in\n let run =\n match runs with\n | Some [ run_id ] -> Munin.Store.find_run store run_id\n | None | Some [] -> Munin.Store.latest_run store ?experiment ()\n | Some _ ->\n Printf.printf \"munin: please specify a single run\\n%!\";\n None\n in\n match run with\n | Some run ->\n let init () = init ~run in\n Mosaic.run { init; update; view; subscriptions }\n | None -> (\n match runs with\n | Some [ run_id ] -> Printf.printf \"munin: run not found: %s\\n%!\" run_id\n | _ ->\n Printf.printf \"munin: no runs found in %s\\n%!\"\n (Munin.Store.root store))\n"
},
{
"path": "packages/munin/lib/tui/munin_tui.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Experiment dashboard TUI.\n\n Terminal-based dashboard for monitoring runs tracked by {!Munin}. Displays\n live metrics, charts, and system resource usage. *)\n\nval run :\n ?root:string -> ?experiment:string -> ?runs:string list -> unit -> unit\n(** [run ()] launches the dashboard.\n\n - [root] defaults to [$RAVEN_TRACKING_DIR] or [$XDG_DATA_HOME/raven/munin].\n - [experiment] filters runs by experiment name.\n - [runs] is a list of run IDs to display; currently exactly one must be\n provided. When omitted, the most recent run is selected automatically. *)\n"
},
{
"path": "packages/munin/lib/tui/overview.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\n\n(* Styles *)\n\nlet section_style =\n Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:12) ~dim:true ()\n\nlet key_style = Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:16) ()\nlet val_style = Ansi.Style.make ~fg:Ansi.Color.white ()\n\n(* Helpers *)\n\nlet is_sys tag = String.length tag > 4 && String.sub tag 0 4 = \"sys/\"\n\nlet format_value (v : Munin.Value.t) =\n match v with\n | `Bool b -> string_of_bool b\n | `Int i -> string_of_int i\n | `Float f ->\n if Float.is_integer f && Float.abs f < 1e9 then Printf.sprintf \"%.0f\" f\n else Printf.sprintf \"%g\" f\n | `String s -> s\n\nlet format_float f =\n if Float.is_integer f && Float.abs f < 1e6 then Printf.sprintf \"%.0f\" f\n else if Float.abs f >= 0.01 && Float.abs f < 1e6 then Printf.sprintf \"%.4f\" f\n else Printf.sprintf \"%.4g\" f\n\n(* View *)\n\nlet section_header label =\n box\n ~size:{ width = pct 100; height = auto }\n [ text ~style:section_style (Printf.sprintf \"\\u{2500}\\u{2500} %s \" label) ]\n\nlet kv_row k v =\n box ~flex_direction:Row ~gap:(gap 1)\n ~size:{ width = pct 100; height = auto }\n [\n text ~style:key_style k;\n box ~flex_grow:1.0 ~size:{ width = auto; height = auto } [];\n text ~style:val_style v;\n ]\n\nlet badge_style =\n Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:10) ~dim:true ()\n\nlet format_summary_value ~(metric_defs : (string * Munin.Run.metric_def) list)\n ~(best_for_tag : string -> float option) k (m : Munin.Run.metric) =\n let def = List.assoc_opt k metric_defs in\n let value =\n match (def, best_for_tag k) with\n | Some { summary = `Min | `Max; _ }, Some bv -> format_float bv\n | _ -> format_float m.value\n in\n let summary_badge =\n match def with\n | Some { summary = `Min; _ } -> \" (min)\"\n | Some { summary = `Max; _ } -> \" (max)\"\n | Some { summary = `Mean; _ } -> \" (mean)\"\n | Some { summary = `None; _ } -> \" (none)\"\n | Some { summary = `Last; _ } | None -> \"\"\n in\n let goal_arrow =\n match def with\n | Some { goal = Some `Minimize; _ } -> \" \\u{2193}\"\n | Some { goal = Some `Maximize; _ } -> \" \\u{2191}\"\n | _ -> \"\"\n in\n box ~flex_direction:Row ~gap:(gap 0)\n ~size:{ width = pct 100; height = auto }\n [\n text ~style:key_style k;\n box ~flex_grow:1.0 ~size:{ width = auto; height = auto } [];\n text ~style:val_style value;\n text ~style:badge_style (summary_badge ^ goal_arrow);\n ]\n\nlet view ~(run : Munin.Run.t)\n ~(latest_metrics : (string * Munin.Run.metric) list)\n ~(step_metrics : string list)\n ~(metric_defs : (string * Munin.Run.metric_def) list)\n ~(best_for_tag : string -> float option) =\n let params = Munin.Run.params run in\n let metrics =\n List.filter\n (fun (k, _) -> (not (is_sys k)) && not (List.mem k step_metrics))\n latest_metrics\n in\n let sections =\n (if params <> [] then\n [ section_header \"Params\" ]\n @ List.map (fun (k, v) -> kv_row k (format_value v)) params\n else [])\n @\n if metrics <> [] then\n [ section_header \"Summary\" ]\n @ List.map\n (fun (k, m) -> format_summary_value ~metric_defs ~best_for_tag k m)\n metrics\n else []\n in\n if sections = [] then box ~size:{ width = px 0; height = px 0 } []\n else\n box ~flex_direction:Column ~gap:(gap 0)\n ~size:{ width = pct 100; height = auto }\n sections\n"
},
{
"path": "packages/munin/lib/tui/system.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\n\n(* Types *)\n\ntype values = {\n cpu_user : float;\n cpu_system : float;\n mem_pct : float;\n mem_gb : float;\n proc_cpu : float;\n proc_mem_mb : float;\n disk_read_mbs : float;\n disk_write_mbs : float;\n disk_util_pct : float;\n}\n\n(* Styles *)\n\nlet label_style = Ansi.Style.make ~fg:Ansi.Color.white ()\nlet bracket_style = Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:10) ()\nlet value_style = Ansi.Style.make ~fg:Ansi.Color.white ()\nlet user_style = Ansi.Style.make ~fg:Ansi.Color.green ()\nlet system_style = Ansi.Style.make ~fg:Ansi.Color.cyan ()\nlet spark_style = Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:14) ()\n\nlet bar_color pct =\n if pct > 80. then Ansi.Color.red\n else if pct > 50. then Ansi.Color.yellow\n else Ansi.Color.green\n\n(* Sparkline rendering *)\n\nlet blocks =\n [|\n \"\\u{2581}\";\n \"\\u{2582}\";\n \"\\u{2583}\";\n \"\\u{2584}\";\n \"\\u{2585}\";\n \"\\u{2586}\";\n \"\\u{2587}\";\n \"\\u{2588}\";\n |]\n\nlet sparkline history ~width =\n let n = List.length history in\n if n = 0 then \"\"\n else\n let values =\n if n <= width then Array.of_list history\n else\n let arr = Array.of_list history in\n Array.sub arr (n - width) width\n in\n let len = Array.length values in\n let lo = Array.fold_left min infinity values in\n let hi = Array.fold_left max neg_infinity values in\n let range = hi -. lo in\n let buf = Buffer.create (len * 3) in\n Array.iter\n (fun v ->\n let idx =\n if range = 0. then 3\n else int_of_float ((v -. lo) /. range *. 7.) |> max 0 |> min 7\n in\n Buffer.add_string buf blocks.(idx))\n values;\n Buffer.contents buf\n\n(* Bar drawing — reused from previous implementation *)\n\nlet draw_bar c ~y ~width ~label ~value_text ~percent ~spark =\n let label_len = String.length label in\n let value_len = String.length value_text in\n Canvas.draw_text c ~x:0 ~y ~text:label ~style:label_style;\n let spark_width = 12 in\n let bar_end = width - spark_width - 2 in\n let bar_start = label_len in\n if bar_end - bar_start < 2 then ()\n else begin\n Canvas.draw_text c ~x:bar_start ~y ~text:\"[\" ~style:bracket_style;\n Canvas.draw_text c ~x:bar_end ~y ~text:\"]\" ~style:bracket_style;\n let inner = bar_end - bar_start - 1 in\n let fill_space = inner - value_len - 1 in\n let fill_count =\n if fill_space > 0 then\n int_of_float (percent /. 100. *. float_of_int fill_space)\n |> max 0 |> min fill_space\n else 0\n in\n let color = bar_color percent in\n let style = Ansi.Style.make ~fg:color () in\n for i = 0 to fill_count - 1 do\n Canvas.draw_text c ~x:(bar_start + 1 + i) ~y ~text:\"|\" ~style\n done;\n let vx = bar_end - value_len in\n if vx > bar_start + 1 then\n Canvas.draw_text c ~x:vx ~y ~text:value_text ~style:value_style;\n Canvas.draw_text c ~x:(bar_end + 2) ~y ~text:spark ~style:spark_style\n end\n\nlet draw_cpu_bar c ~y ~width ~label ~cpu_user ~cpu_system ~spark =\n let label_len = String.length label in\n let total = cpu_user +. cpu_system in\n let value_text = Printf.sprintf \"%.0f%%\" total in\n let value_len = String.length value_text in\n Canvas.draw_text c ~x:0 ~y ~text:label ~style:label_style;\n let spark_width = 12 in\n let bar_end = width - spark_width - 2 in\n let bar_start = label_len in\n if bar_end - bar_start < 2 then ()\n else begin\n Canvas.draw_text c ~x:bar_start ~y ~text:\"[\" ~style:bracket_style;\n Canvas.draw_text c ~x:bar_end ~y ~text:\"]\" ~style:bracket_style;\n let inner = bar_end - bar_start - 1 in\n let fill_space = inner - value_len - 1 in\n let user_count =\n if fill_space > 0 then\n int_of_float (cpu_user /. 100. *. float_of_int fill_space)\n |> max 0 |> min fill_space\n else 0\n in\n let system_count =\n if fill_space > 0 then\n int_of_float (cpu_system /. 100. *. float_of_int fill_space)\n |> max 0\n |> min (fill_space - user_count)\n else 0\n in\n for i = 0 to user_count - 1 do\n Canvas.draw_text c ~x:(bar_start + 1 + i) ~y ~text:\"|\" ~style:user_style\n done;\n for i = 0 to system_count - 1 do\n Canvas.draw_text c\n ~x:(bar_start + 1 + user_count + i)\n ~y ~text:\"|\" ~style:system_style\n done;\n let vx = bar_end - value_len in\n if vx > bar_start + 1 then\n Canvas.draw_text c ~x:vx ~y ~text:value_text ~style:value_style;\n Canvas.draw_text c ~x:(bar_end + 2) ~y ~text:spark ~style:spark_style\n end\n\nlet draw_value_line c ~y ~width ~label ~value_text ~spark =\n let label_len = String.length label in\n Canvas.draw_text c ~x:0 ~y ~text:label ~style:label_style;\n let spark_width = 12 in\n let vx = width - spark_width - 2 - String.length value_text in\n if vx > label_len then\n Canvas.draw_text c ~x:vx ~y ~text:value_text ~style:value_style;\n Canvas.draw_text c ~x:(width - spark_width) ~y ~text:spark ~style:spark_style\n\n(* View *)\n\nlet extract_values history = List.map snd history\n\nlet view (v : values) ~(history_for_tag : string -> (int * float) list) =\n let spark tag = sparkline (extract_values (history_for_tag tag)) ~width:10 in\n let cpu_spark =\n let h1 = extract_values (history_for_tag \"sys/cpu_user\") in\n let h2 = extract_values (history_for_tag \"sys/cpu_system\") in\n let rec zip_add a b =\n match (a, b) with\n | x :: xs, y :: ys -> (x +. y) :: zip_add xs ys\n | [], _ | _, [] -> []\n in\n sparkline (zip_add h1 h2) ~width:10\n in\n let format_rate mbs =\n if mbs >= 1024. then Printf.sprintf \"%.1f GB/s\" (mbs /. 1024.)\n else if mbs >= 1. then Printf.sprintf \"%.1f MB/s\" mbs\n else Printf.sprintf \"%.0f KB/s\" (mbs *. 1024.)\n in\n canvas\n ~size:{ width = pct 100; height = px 7 }\n (fun c ~delta:_ ->\n Canvas.clear c;\n let w = Canvas.width c in\n draw_cpu_bar c ~y:0 ~width:w ~label:\"CPU \" ~cpu_user:v.cpu_user\n ~cpu_system:v.cpu_system ~spark:cpu_spark;\n draw_bar c ~y:1 ~width:w ~label:\"Mem \"\n ~value_text:(Printf.sprintf \"%.0f%% %.1fG\" v.mem_pct v.mem_gb)\n ~percent:v.mem_pct ~spark:(spark \"sys/mem_used_pct\");\n draw_bar c ~y:2 ~width:w ~label:\"Proc \"\n ~value_text:(Printf.sprintf \"%.0f%%\" v.proc_cpu)\n ~percent:v.proc_cpu ~spark:(spark \"sys/proc_cpu_pct\");\n draw_value_line c ~y:3 ~width:w ~label:\"RSS \"\n ~value_text:(Printf.sprintf \"%.0fM\" v.proc_mem_mb)\n ~spark:(spark \"sys/proc_mem_mb\");\n draw_bar c ~y:4 ~width:w ~label:\"Disk \"\n ~value_text:(Printf.sprintf \"%.0f%%\" v.disk_util_pct)\n ~percent:v.disk_util_pct\n ~spark:(spark \"sys/disk_util_pct\");\n draw_value_line c ~y:5 ~width:w ~label:\" R \"\n ~value_text:(format_rate v.disk_read_mbs)\n ~spark:(spark \"sys/disk_read_mbs\");\n draw_value_line c ~y:6 ~width:w ~label:\" W \"\n ~value_text:(format_rate v.disk_write_mbs)\n ~spark:(spark \"sys/disk_write_mbs\"))\n"
},
{
"path": "packages/munin/lib/tui/theme.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Mosaic\nmodule Charts = Matrix_charts\n\n(* Run status *)\n\ntype run_status = Live | Stopped | Done | Failed | Killed\n\nlet status_label = function\n | Live -> \"LIVE\"\n | Stopped -> \"Stopped\"\n | Done -> \"Done\"\n | Failed -> \"Failed\"\n | Killed -> \"Killed\"\n\nlet status_color = function\n | Live -> Ansi.Color.green\n | Stopped -> Ansi.Color.grayscale ~level:12\n | Done -> Ansi.Color.of_rgb 80 140 200\n | Failed -> Ansi.Color.red\n | Killed -> Ansi.Color.yellow\n\n(* Smooth level *)\n\ntype smooth = Off | Light | Medium | Heavy\n\nlet smooth_alpha = function\n | Off -> None\n | Light -> Some 0.5\n | Medium -> Some 0.3\n | Heavy -> Some 0.15\n\nlet next_smooth = function\n | Off -> Light\n | Light -> Medium\n | Medium -> Heavy\n | Heavy -> Off\n\nlet smooth_display = function Off -> 0 | Light -> 1 | Medium -> 2 | Heavy -> 3\n\n(* Shared styles *)\n\nlet muted_style = Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:14) ()\n\nlet axis_style =\n Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:12) ~dim:true ()\n\nlet grid_style =\n Ansi.Style.make ~fg:(Ansi.Color.grayscale ~level:6) ~dim:true ()\n\nlet header_bg = Ansi.Color.grayscale ~level:2\n\n(* Shared helpers *)\n\nlet last_value history =\n let rec go = function\n | [] -> None\n | [ (_, v) ] -> Some v\n | _ :: rest -> go rest\n in\n go history\n\n(* Chart drawing *)\n\nlet draw_metric_chart ~compact history grid ~width ~height =\n if history = [] then ()\n else\n let data =\n Array.of_list\n (List.map (fun (step, value) -> (float_of_int step, value)) history)\n in\n let margins, x_ticks, y_ticks, y_format =\n if compact then ((1, 0, 0, 2), 4, 2, fun _ v -> Printf.sprintf \"%.1f\" v)\n else ((1, 1, 1, 4), 6, 4, fun _ v -> Printf.sprintf \"%.4g\" v)\n in\n let chart =\n Charts.empty ()\n |> Charts.with_frame (Charts.manual_frame ~margins ())\n |> Charts.with_axes\n ~x:\n (Charts.Axis.default\n |> Charts.Axis.with_ticks x_ticks\n |> Charts.Axis.with_style axis_style)\n ~y:\n (Charts.Axis.default\n |> Charts.Axis.with_ticks y_ticks\n |> Charts.Axis.with_style axis_style\n |> Charts.Axis.with_format y_format)\n |> Charts.with_grid\n (Charts.Gridlines.default\n |> Charts.Gridlines.with_style grid_style\n |> Charts.Gridlines.with_x true\n |> Charts.Gridlines.with_y true)\n |> Charts.line ~id:\"metric\" ~resolution:`Braille2x4\n ~style:(Ansi.Style.make ~fg:Ansi.Color.cyan ())\n ~x:fst ~y:snd data\n in\n ignore (Charts.draw chart grid ~width ~height)\n"
},
{
"path": "packages/munin/lib/value.ml",
"content": "type t = [ `Bool of bool | `Int of int | `Float of float | `String of string ]\n\nlet pp ppf = function\n | `Bool b -> Format.pp_print_bool ppf b\n | `Int n -> Format.pp_print_int ppf n\n | `Float f -> Format.fprintf ppf \"%g\" f\n | `String s -> Format.fprintf ppf \"%S\" s\n\nlet to_float = function\n | `Float f -> Some f\n | `Int n -> Some (Float.of_int n)\n | _ -> None\n\nlet to_int = function\n | `Int n -> Some n\n | `Float f when Float.is_integer f -> Some (Float.to_int f)\n | _ -> None\n\nlet to_string = function `String s -> Some s | _ -> None\nlet to_bool = function `Bool b -> Some b | _ -> None\n\nlet to_json : t -> Jsont.json = function\n | `Bool b -> Jsont.Json.bool b\n | `Int n -> Jsont.Json.int n\n | `Float f -> Jsont.Json.number f\n | `String s -> Jsont.Json.string s\n\nlet of_json : Jsont.json -> t = function\n | Jsont.Bool (b, _) -> `Bool b\n | Jsont.Number (f, _) ->\n if Float.is_integer f && Float.abs f < 4503599627370496. then\n `Int (Float.to_int f)\n else `Float f\n | Jsont.String (s, _) -> `String s\n | _ -> `String \"\"\n"
},
{
"path": "packages/munin/lib/value.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** Scalar values for parameters, summaries, and metadata. *)\n\ntype t = [ `Bool of bool | `Int of int | `Float of float | `String of string ]\n(** The type for scalar values. *)\n\nval pp : Format.formatter -> t -> unit\n(** [pp ppf v] pretty-prints [v]. *)\n\nval to_float : t -> float option\n(** [to_float v] extracts a float. [`Int] values are promoted. *)\n\nval to_int : t -> int option\n(** [to_int v] extracts an int. [`Float] values are truncated if integral. *)\n\nval to_string : t -> string option\n(** [to_string v] extracts a string. *)\n\nval to_bool : t -> bool option\n(** [to_bool v] extracts a bool. *)\n\n(**/**)\n\nval to_json : t -> Jsont.json\nval of_json : Jsont.json -> t\n\n(**/**)\n"
},
{
"path": "packages/munin/test/dune",
"content": "(test\n (name test_munin)\n (package munin)\n (libraries munin munin.sys unix threads.posix windtrap jsont))\n"
},
{
"path": "packages/munin/test/test_munin.ml",
"content": "open Windtrap\n\n(* Helpers *)\n\nlet with_temp_dir f =\n let base = Filename.temp_file \"munin\" \"test\" in\n Sys.remove base;\n Unix.mkdir base 0o755;\n Fun.protect\n ~finally:(fun () -> ignore (Sys.command (\"rm -rf \" ^ Filename.quote base)))\n (fun () -> f base)\n\nlet write_text path text =\n let oc = open_out path in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc text)\n\nlet make_file root name text =\n let path = Filename.concat root name in\n write_text path text;\n path\n\nlet make_dir root name files =\n let dir = Filename.concat root name in\n Unix.mkdir dir 0o755;\n List.iter\n (fun (fname, text) -> write_text (Filename.concat dir fname) text)\n files;\n dir\n\n(* Session lifecycle *)\n\nlet test_start_creates_run_dir () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n let run = Session.run session in\n is_true ~msg:\"run dir exists\" (Sys.file_exists (Run.dir run));\n is_true ~msg:\"manifest exists\"\n (Sys.file_exists (Filename.concat (Run.dir run) \"run.json\"));\n is_true ~msg:\"status is running\" (Run.status run = `running);\n Session.finish session ()\n\nlet test_finish_sets_status () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n let run = Session.run session in\n is_true ~msg:\"status finished\" (Run.status run = `finished);\n is_true ~msg:\"ended_at set\" (Option.is_some (Run.ended_at run))\n\nlet test_finish_failed () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish ~status:`failed session ();\n is_true ~msg:\"status failed\" (Run.status (Session.run session) = `failed)\n\nlet test_finish_killed () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish ~status:`killed session ();\n is_true ~msg:\"status killed\" (Run.status (Session.run session) = `killed)\n\nlet test_finish_idempotent () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n Session.finish ~status:`failed session ();\n is_true ~msg:\"still finished\" (Run.status (Session.run session) = `finished)\n\nlet test_with_run_success () =\n with_temp_dir @@ fun root ->\n let result =\n Session.with_run ~root ~experiment:\"exp\" (fun session ->\n Session.log_metric session ~step:1 \"x\" 1.0;\n 42)\n in\n equal ~msg:\"return value\" int 42 result;\n let store = Store.open_ ~root () in\n let run = Option.get (Store.latest_run store ()) in\n is_true ~msg:\"status finished\" (Run.status run = `finished)\n\nlet test_with_run_exception () =\n with_temp_dir @@ fun root ->\n let raised =\n try\n ignore\n (Session.with_run ~root ~experiment:\"exp\" (fun _session ->\n failwith \"boom\"));\n false\n with Failure _ -> true\n in\n is_true ~msg:\"exception re-raised\" raised;\n let store = Store.open_ ~root () in\n let run = Option.get (Store.latest_run store ()) in\n is_true ~msg:\"status failed\" (Run.status run = `failed)\n\nlet test_resume () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 1.0;\n let run = Session.run session in\n is_true ~msg:\"resumable before finish\" (Run.resumable run);\n let resumed = Session.resume run in\n Session.log_metric resumed ~step:2 \"x\" 0.5;\n Session.finish resumed ();\n let final = Session.run resumed in\n equal ~msg:\"history length\" int 2 (List.length (Run.metric_history final \"x\"));\n is_false ~msg:\"not resumable after finish\" (Run.resumable final)\n\nlet test_resume_finished_raises () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n let run = Session.run session in\n raises_invalid_arg \"Munin.Session.resume: run is not resumable\" (fun () ->\n ignore (Session.resume run))\n\nlet test_ops_after_finish_ignored () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 1.0;\n Session.finish session ();\n Session.log_metric session ~step:2 \"x\" 2.0;\n Session.set_notes session (Some \"late\");\n Session.add_tags session [ \"late\" ];\n Session.set_summary session [ (\"late\", `Float 1.0) ];\n let run = Session.run session in\n equal ~msg:\"only one metric\" int 1 (List.length (Run.metric_history run \"x\"));\n is_true ~msg:\"no late note\" (Run.notes run = None);\n equal ~msg:\"no tags\" (list string) [] (Run.tags run)\n\nlet lifecycle =\n [\n test \"start creates run dir\" test_start_creates_run_dir;\n test \"finish sets status\" test_finish_sets_status;\n test \"finish failed\" test_finish_failed;\n test \"finish killed\" test_finish_killed;\n test \"finish is idempotent\" test_finish_idempotent;\n test \"with_run success\" test_with_run_success;\n test \"with_run exception\" test_with_run_exception;\n test \"resume\" test_resume;\n test \"resume finished raises\" test_resume_finished_raises;\n test \"ops after finish ignored\" test_ops_after_finish_ignored;\n ]\n\n(* Scalars *)\n\nlet test_log_metric () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"train/loss\" 1.5;\n Session.finish session ();\n let run = Session.run session in\n let m =\n match List.assoc_opt \"train/loss\" (Run.latest_metrics run) with\n | Some m -> m\n | None -> failwith \"missing\"\n in\n equal ~msg:\"step\" int 1 m.step;\n equal ~msg:\"value\" (float 0.0) 1.5 m.value\n\nlet test_log_metrics_batch () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metrics session ~step:1 [ (\"train/loss\", 1.0); (\"val/acc\", 0.6) ];\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"two keys\" int 2 (List.length (Run.metric_keys run))\n\nlet test_metric_history_chronological () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 3.0;\n Session.log_metric session ~step:2 \"x\" 2.0;\n Session.log_metric session ~step:3 \"x\" 1.0;\n Session.finish session ();\n let run = Session.run session in\n let history = Run.metric_history run \"x\" in\n equal ~msg:\"length\" int 3 (List.length history);\n let values = List.map (fun (m : Run.metric) -> m.value) history in\n equal ~msg:\"order\" (list (float 0.0)) [ 3.0; 2.0; 1.0 ] values\n\nlet test_latest_metrics () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 1.0;\n Session.log_metric session ~step:2 \"x\" 2.0;\n Session.finish session ();\n let run = Session.run session in\n let latest =\n match List.assoc_opt \"x\" (Run.latest_metrics run) with\n | Some m -> m\n | None -> failwith \"missing\"\n in\n equal ~msg:\"latest step\" int 2 latest.step;\n equal ~msg:\"latest value\" (float 0.0) 2.0 latest.value\n\nlet test_metric_keys_sorted () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"z/loss\" 1.0;\n Session.log_metric session ~step:1 \"a/acc\" 0.5;\n Session.log_metric session ~step:1 \"m/lr\" 0.01;\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"sorted keys\" (list string)\n [ \"a/acc\"; \"m/lr\"; \"z/loss\" ]\n (Run.metric_keys run)\n\nlet test_explicit_timestamp () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 ~timestamp:42.0 \"x\" 1.0;\n Session.finish session ();\n let run = Session.run session in\n let m =\n match List.assoc_opt \"x\" (Run.latest_metrics run) with\n | Some m -> m\n | None -> failwith \"missing\"\n in\n equal ~msg:\"timestamp\" (float 0.0) 42.0 m.timestamp\n\nlet test_missing_metric_history_empty () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"empty history\" (list int) []\n (List.map\n (fun (m : Run.metric) -> m.step)\n (Run.metric_history run \"nonexistent\"))\n\nlet scalars =\n [\n test \"log_metric\" test_log_metric;\n test \"log_metrics batch\" test_log_metrics_batch;\n test \"history chronological\" test_metric_history_chronological;\n test \"latest_metrics\" test_latest_metrics;\n test \"metric_keys sorted\" test_metric_keys_sorted;\n test \"explicit timestamp\" test_explicit_timestamp;\n test \"missing key history empty\" test_missing_metric_history_empty;\n ]\n\n(* Metric definitions *)\n\nlet test_define_metric_summary_and_goal () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"train/loss\" ~summary:`Min ~goal:`Minimize ();\n Session.define_metric session \"val/acc\" ~summary:`Max ~goal:`Maximize ();\n Session.finish session ();\n let run = Session.run session in\n let defs = Run.metric_defs run in\n equal ~msg:\"two defs\" int 2 (List.length defs);\n let loss_def = List.assoc \"train/loss\" defs in\n is_true ~msg:\"loss summary min\" (loss_def.summary = `Min);\n is_true ~msg:\"loss goal minimize\" (loss_def.goal = Some `Minimize);\n let acc_def = List.assoc \"val/acc\" defs in\n is_true ~msg:\"acc summary max\" (acc_def.summary = `Max);\n is_true ~msg:\"acc goal maximize\" (acc_def.goal = Some `Maximize)\n\nlet test_define_metric_all_summaries () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"a\" ~summary:`Min ();\n Session.define_metric session \"b\" ~summary:`Max ();\n Session.define_metric session \"c\" ~summary:`Mean ();\n Session.define_metric session \"d\" ~summary:`Last ();\n Session.define_metric session \"e\" ~summary:`None ();\n Session.finish session ();\n let run = Session.run session in\n let defs = Run.metric_defs run in\n equal ~msg:\"five defs\" int 5 (List.length defs);\n is_true ~msg:\"a=Min\" ((List.assoc \"a\" defs).summary = `Min);\n is_true ~msg:\"b=Max\" ((List.assoc \"b\" defs).summary = `Max);\n is_true ~msg:\"c=Mean\" ((List.assoc \"c\" defs).summary = `Mean);\n is_true ~msg:\"d=Last\" ((List.assoc \"d\" defs).summary = `Last);\n is_true ~msg:\"e=None\" ((List.assoc \"e\" defs).summary = `None)\n\nlet test_define_metric_step_metric () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"train/loss\" ~step_metric:\"epoch\" ();\n Session.finish session ();\n let run = Session.run session in\n let def = List.assoc \"train/loss\" (Run.metric_defs run) in\n equal ~msg:\"step_metric\" (option string) (Some \"epoch\") def.step_metric\n\nlet test_define_metric_default_summary () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"x\" ();\n Session.finish session ();\n let run = Session.run session in\n let def = List.assoc \"x\" (Run.metric_defs run) in\n is_true ~msg:\"default summary is Last\" (def.summary = `Last)\n\nlet test_metric_defs_sorted () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"z\" ();\n Session.define_metric session \"a\" ();\n Session.define_metric session \"m\" ();\n Session.finish session ();\n let run = Session.run session in\n let keys = List.map fst (Run.metric_defs run) in\n equal ~msg:\"sorted\" (list string) [ \"a\"; \"m\"; \"z\" ] keys\n\nlet metric_definitions =\n [\n test \"summary and goal\" test_define_metric_summary_and_goal;\n test \"all summary modes\" test_define_metric_all_summaries;\n test \"step_metric\" test_define_metric_step_metric;\n test \"default summary is Last\" test_define_metric_default_summary;\n test \"defs sorted\" test_metric_defs_sorted;\n ]\n\n(* Metadata *)\n\nlet test_set_notes () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.set_notes session (Some \"hello\");\n Session.finish session ();\n equal ~msg:\"notes set\" (option string) (Some \"hello\")\n (Run.notes (Session.run session))\n\nlet test_set_notes_replace () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.set_notes session (Some \"first\");\n Session.set_notes session (Some \"second\");\n Session.finish session ();\n equal ~msg:\"notes replaced\" (option string) (Some \"second\")\n (Run.notes (Session.run session))\n\nlet test_clear_notes () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" ~notes:\"initial\" () in\n Session.set_notes session None;\n Session.finish session ();\n equal ~msg:\"notes cleared\" (option string) None\n (Run.notes (Session.run session))\n\nlet test_set_summary_merge () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.set_summary session [ (\"a\", `Float 1.0) ];\n Session.set_summary session [ (\"b\", `Float 2.0); (\"a\", `Float 3.0) ];\n Session.finish session ();\n let run = Session.run session in\n let summary = Run.summary run in\n equal ~msg:\"two keys\" int 2 (List.length summary);\n is_true ~msg:\"a replaced\"\n (Value.to_float (Option.get (Run.find_summary run \"a\")) = Some 3.0);\n is_true ~msg:\"b present\"\n (Value.to_float (Option.get (Run.find_summary run \"b\")) = Some 2.0)\n\nlet test_find_summary_missing () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n equal ~msg:\"missing key\" (option unit) None\n (Option.map ignore (Run.find_summary (Session.run session) \"nope\"))\n\nlet test_add_tags () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" ~tags:[ \"initial\" ] () in\n Session.add_tags session [ \"added\" ];\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"tags\" (list string) [ \"initial\"; \"added\" ] (Run.tags run)\n\nlet test_add_tags_dedup () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" ~tags:[ \"a\" ] () in\n Session.add_tags session [ \"a\"; \"b\" ];\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"deduped\" (list string) [ \"a\"; \"b\" ] (Run.tags run)\n\nlet test_add_tags_empty_noop () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.add_tags session [];\n Session.finish session ();\n equal ~msg:\"no tags\" (list string) [] (Run.tags (Session.run session))\n\nlet test_params_immutable () =\n with_temp_dir @@ fun root ->\n let session =\n Session.start ~root ~experiment:\"exp\"\n ~params:[ (\"lr\", `Float 0.001); (\"bs\", `Int 32) ]\n ()\n in\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"param count\" int 2 (List.length (Run.params run))\n\nlet metadata =\n [\n test \"set_notes\" test_set_notes;\n test \"set_notes replace\" test_set_notes_replace;\n test \"clear notes\" test_clear_notes;\n test \"set_summary merge\" test_set_summary_merge;\n test \"find_summary missing\" test_find_summary_missing;\n test \"add_tags\" test_add_tags;\n test \"add_tags dedup\" test_add_tags_dedup;\n test \"add_tags empty noop\" test_add_tags_empty_noop;\n test \"params immutable\" test_params_immutable;\n ]\n\n(* Provenance *)\n\nlet test_provenance_fields () =\n with_temp_dir @@ fun root ->\n let session =\n Session.start ~root ~experiment:\"exp\" ~command:[ \"ocaml\"; \"train.ml\" ]\n ~cwd:root ~hostname:\"node0\" ~pid:42 ~git_commit:\"abc123\" ~git_dirty:true\n ()\n in\n Session.finish session ();\n let prov = Run.provenance (Session.run session) in\n equal ~msg:\"command\" (list string) [ \"ocaml\"; \"train.ml\" ] prov.command;\n equal ~msg:\"cwd\" string root prov.cwd;\n equal ~msg:\"hostname\" (option string) (Some \"node0\") prov.hostname;\n equal ~msg:\"pid\" int 42 prov.pid;\n equal ~msg:\"git_commit\" (option string) (Some \"abc123\") prov.git_commit;\n equal ~msg:\"git_dirty\" (option bool) (Some true) prov.git_dirty\n\nlet test_capture_env () =\n with_temp_dir @@ fun root ->\n Unix.putenv \"MUNIN_TEST_A\" \"alpha\";\n let session =\n Session.start ~root ~experiment:\"exp\" ~capture_env:[ \"MUNIN_TEST_A\" ] ()\n in\n Session.finish session ();\n let prov = Run.provenance (Session.run session) in\n equal ~msg:\"env captured\"\n (list (pair string string))\n [ (\"MUNIN_TEST_A\", \"alpha\") ]\n prov.env\n\nlet test_capture_env_missing () =\n with_temp_dir @@ fun root ->\n (* Use a variable name that won't exist *)\n let session =\n Session.start ~root ~experiment:\"exp\"\n ~capture_env:[ \"MUNIN_NONEXISTENT_9999\" ]\n ()\n in\n Session.finish session ();\n let prov = Run.provenance (Session.run session) in\n equal ~msg:\"missing env omitted\" (list (pair string string)) [] prov.env\n\nlet test_explicit_env () =\n with_temp_dir @@ fun root ->\n let session =\n Session.start ~root ~experiment:\"exp\" ~env:[ (\"KEY\", \"value\") ] ()\n in\n Session.finish session ();\n let prov = Run.provenance (Session.run session) in\n equal ~msg:\"explicit env\"\n (list (pair string string))\n [ (\"KEY\", \"value\") ]\n prov.env\n\nlet provenance_tests =\n [\n test \"all fields round-trip\" test_provenance_fields;\n test \"capture_env\" test_capture_env;\n test \"capture_env missing var\" test_capture_env_missing;\n test \"explicit env\" test_explicit_env;\n ]\n\n(* Run loading *)\n\nlet test_run_load () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n let id = Run.id (Session.run session) in\n is_some ~msg:\"load existing\" (Run.load ~root ~experiment:\"exp\" ~id)\n\nlet test_run_load_missing () =\n with_temp_dir @@ fun root ->\n let _store = Store.open_ ~root () in\n is_none ~msg:\"load missing\"\n (Run.load ~root ~experiment:\"exp\" ~id:\"nonexistent\")\n\nlet test_run_list_sorted_descending () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp\" () in\n Session.finish s2 ();\n let runs = Run.list ~root ~experiment:\"exp\" () in\n equal ~msg:\"count\" int 2 (List.length runs);\n let ids = List.map Run.id runs in\n (* Sorted descending by id *)\n is_true ~msg:\"descending order\"\n (String.compare (List.nth ids 0) (List.nth ids 1) > 0)\n\nlet test_run_list_filter_status () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp\" () in\n Session.finish ~status:`failed s2 ();\n equal ~msg:\"finished only\" int 1\n (List.length (Run.list ~root ~experiment:\"exp\" ~status:`finished ()));\n equal ~msg:\"failed only\" int 1\n (List.length (Run.list ~root ~experiment:\"exp\" ~status:`failed ()))\n\nlet test_run_list_filter_tag () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp\" ~tags:[ \"train\" ] () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp\" ~tags:[ \"eval\" ] () in\n Session.finish s2 ();\n equal ~msg:\"train tagged\" int 1\n (List.length (Run.list ~root ~experiment:\"exp\" ~tag:\"train\" ()));\n equal ~msg:\"eval tagged\" int 1\n (List.length (Run.list ~root ~experiment:\"exp\" ~tag:\"eval\" ()))\n\nlet test_run_list_filter_parent () =\n with_temp_dir @@ fun root ->\n let parent = Session.start ~root ~experiment:\"exp\" () in\n Session.finish parent ();\n let parent_run = Session.run parent in\n let child = Session.start ~root ~experiment:\"exp\" ~parent:parent_run () in\n Session.finish child ();\n let _other = Session.start ~root ~experiment:\"exp\" () in\n Session.finish _other ();\n equal ~msg:\"one child\" int 1\n (List.length\n (Run.list ~root ~experiment:\"exp\" ~parent:(Run.id parent_run) ()))\n\nlet test_run_children () =\n with_temp_dir @@ fun root ->\n let parent = Session.start ~root ~experiment:\"exp\" () in\n Session.finish parent ();\n let parent_run = Session.run parent in\n let c1 = Session.start ~root ~experiment:\"exp\" ~parent:parent_run () in\n Session.finish c1 ();\n let c2 = Session.start ~root ~experiment:\"exp\" ~parent:parent_run () in\n Session.finish c2 ();\n let children = Run.children parent_run in\n equal ~msg:\"two children\" int 2 (List.length children);\n List.iter\n (fun child ->\n equal ~msg:\"parent_id\" (option string)\n (Some (Run.id parent_run))\n (Run.parent_id child))\n children\n\nlet test_run_name () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp\" ~name:\"baseline\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp\" () in\n Session.finish s2 ();\n equal ~msg:\"named\" (option string) (Some \"baseline\")\n (Run.name (Session.run s1));\n equal ~msg:\"unnamed\" (option string) None (Run.name (Session.run s2))\n\nlet test_experiment_name () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"my-exp\" () in\n Session.finish session ();\n equal ~msg:\"experiment name\" string \"my-exp\"\n (Run.experiment_name (Session.run session))\n\nlet run_loading =\n [\n test \"load existing\" test_run_load;\n test \"load missing\" test_run_load_missing;\n test \"list sorted descending\" test_run_list_sorted_descending;\n test \"list filter status\" test_run_list_filter_status;\n test \"list filter tag\" test_run_list_filter_tag;\n test \"list filter parent\" test_run_list_filter_parent;\n test \"children\" test_run_children;\n test \"name\" test_run_name;\n test \"experiment_name\" test_experiment_name;\n ]\n\n(* Artifacts *)\n\nlet test_file_artifact () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"weights.bin\" \"model weights\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n let artifact =\n Session.log_artifact session ~name:\"model\" ~kind:`checkpoint ~path:src ()\n in\n Session.finish session ();\n equal ~msg:\"name\" string \"model\" (Artifact.name artifact);\n equal ~msg:\"version\" string \"v1\" (Artifact.version artifact);\n is_true ~msg:\"kind\" (Artifact.kind artifact = `checkpoint);\n is_true ~msg:\"payload file\" (Artifact.payload artifact = `file);\n is_true ~msg:\"size positive\" (Artifact.size_bytes artifact > 0);\n is_true ~msg:\"path exists\" (Sys.file_exists (Artifact.path artifact));\n is_true ~msg:\"digest is sha256 hex\"\n (String.length (Artifact.digest artifact) = 64)\n\nlet test_dir_artifact () =\n with_temp_dir @@ fun root ->\n let dir =\n make_dir root \"dataset\" [ (\"a.txt\", \"data a\"); (\"b.txt\", \"data b\") ]\n in\n let session = Session.start ~root ~experiment:\"exp\" () in\n let artifact =\n Session.log_artifact session ~name:\"data\" ~kind:`dataset ~path:dir ()\n in\n Session.finish session ();\n is_true ~msg:\"payload dir\" (Artifact.payload artifact = `dir);\n is_true ~msg:\"kind dataset\" (Artifact.kind artifact = `dataset);\n is_true ~msg:\"blob dir exists\" (Sys.file_exists (Artifact.path artifact))\n\nlet test_artifact_versioning () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"v1 weights\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n let a1 =\n Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src ()\n in\n write_text src \"v2 weights\";\n let a2 =\n Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src ()\n in\n Session.finish session ();\n equal ~msg:\"first version\" string \"v1\" (Artifact.version a1);\n equal ~msg:\"second version\" string \"v2\" (Artifact.version a2)\n\nlet test_artifact_aliases () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"weights\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n let artifact =\n Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src\n ~aliases:[ \"best\"; \"latest\" ] ()\n in\n Session.finish session ();\n is_true ~msg:\"has best alias\" (Artifact.has_alias artifact \"best\");\n is_true ~msg:\"has latest alias\" (Artifact.has_alias artifact \"latest\");\n is_false ~msg:\"no nope alias\" (Artifact.has_alias artifact \"nope\")\n\nlet test_artifact_alias_resolution () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"weights\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n ignore\n (Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src\n ~aliases:[ \"latest\" ] ());\n write_text src \"v2 weights\";\n ignore\n (Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src\n ~aliases:[ \"latest\" ] ());\n Session.finish session ();\n let resolved = Artifact.load ~root ~name:\"model\" ~version:\"latest\" in\n is_true ~msg:\"alias resolves to v2\"\n (match resolved with Some a -> Artifact.version a = \"v2\" | None -> false);\n let v1 = Artifact.load ~root ~name:\"model\" ~version:\"v1\" in\n is_some ~msg:\"explicit v1 loads\" v1\n\nlet test_artifact_metadata () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"weights\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n let artifact =\n Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src\n ~metadata:[ (\"framework\", `String \"rune\") ]\n ()\n in\n Session.finish session ();\n equal ~msg:\"metadata count\" int 1 (List.length (Artifact.metadata artifact))\n\nlet test_artifact_lineage () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"weights\" in\n let producer = Session.start ~root ~experiment:\"exp\" ~name:\"producer\" () in\n let artifact =\n Session.log_artifact producer ~name:\"model\" ~kind:`model ~path:src ()\n in\n let consumer = Session.start ~root ~experiment:\"exp\" ~name:\"consumer\" () in\n Session.use_artifact consumer artifact;\n Session.finish producer ();\n Session.finish consumer ();\n let producer_run = Session.run producer in\n let consumer_run = Session.run consumer in\n equal ~msg:\"producer_run_id\" (option string)\n (Some (Run.id producer_run))\n (Artifact.producer_run_id artifact);\n (* Reload to see consumer linkage *)\n let reloaded =\n match Artifact.load ~root ~name:\"model\" ~version:\"v1\" with\n | Some a -> a\n | None -> failwith \"missing artifact\"\n in\n equal ~msg:\"consumer_run_ids\" (list string)\n [ Run.id consumer_run ]\n (Artifact.consumer_run_ids reloaded);\n equal ~msg:\"output artifacts\" int 1\n (List.length (Run.output_artifacts producer_run));\n equal ~msg:\"input artifacts\" int 1\n (List.length (Run.input_artifacts consumer_run))\n\nlet test_log_artifact_closed_raises () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"weights\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n raises_failure \"Munin.Session.log_artifact: closed session\" (fun () ->\n ignore\n (Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src ()))\n\nlet test_log_artifact_missing_path_raises () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n raises_invalid_arg\n \"Munin.Session.log_artifact: path does not exist: /nonexistent/path\"\n (fun () ->\n ignore\n (Session.log_artifact session ~name:\"model\" ~kind:`model\n ~path:\"/nonexistent/path\" ()))\n\nlet test_artifact_content_dedup () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"same content\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n let a1 =\n Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src ()\n in\n let a2 =\n Session.log_artifact session ~name:\"backup\" ~kind:`model ~path:src ()\n in\n Session.finish session ();\n equal ~msg:\"same digest\" string (Artifact.digest a1) (Artifact.digest a2)\n\nlet artifacts =\n [\n test \"file artifact\" test_file_artifact;\n test \"directory artifact\" test_dir_artifact;\n test \"versioning\" test_artifact_versioning;\n test \"aliases\" test_artifact_aliases;\n test \"alias resolution\" test_artifact_alias_resolution;\n test \"metadata\" test_artifact_metadata;\n test \"lineage\" test_artifact_lineage;\n test \"closed session raises\" test_log_artifact_closed_raises;\n test \"missing path raises\" test_log_artifact_missing_path_raises;\n test \"content dedup\" test_artifact_content_dedup;\n ]\n\n(* Store *)\n\nlet test_store_open () =\n with_temp_dir @@ fun root ->\n let store = Store.open_ ~root () in\n equal ~msg:\"root\" string root (Store.root store);\n is_true ~msg:\"experiments dir exists\"\n (Sys.file_exists (Filename.concat root \"experiments\"));\n is_true ~msg:\"artifacts dir exists\"\n (Sys.file_exists (Filename.concat root \"artifacts\"))\n\nlet test_store_list_experiments () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"mnist\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"cifar\" () in\n Session.finish s2 ();\n let store = Store.open_ ~root () in\n let exps = Store.list_experiments store in\n equal ~msg:\"two experiments\" int 2 (List.length exps);\n is_true ~msg:\"has mnist\" (List.mem \"mnist\" exps);\n is_true ~msg:\"has cifar\" (List.mem \"cifar\" exps)\n\nlet test_store_list_runs_all () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp1\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp2\" () in\n Session.finish s2 ();\n let store = Store.open_ ~root () in\n equal ~msg:\"all runs\" int 2 (List.length (Store.list_runs store ()))\n\nlet test_store_list_runs_experiment () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp1\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp2\" () in\n Session.finish s2 ();\n let store = Store.open_ ~root () in\n equal ~msg:\"exp1 runs\" int 1\n (List.length (Store.list_runs store ~experiment:\"exp1\" ()));\n equal ~msg:\"exp2 runs\" int 1\n (List.length (Store.list_runs store ~experiment:\"exp2\" ()))\n\nlet test_store_find_run () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n let id = Run.id (Session.run session) in\n let store = Store.open_ ~root () in\n is_some ~msg:\"found\" (Store.find_run store id)\n\nlet test_store_find_run_missing () =\n with_temp_dir @@ fun root ->\n let store = Store.open_ ~root () in\n is_none ~msg:\"not found\" (Store.find_run store \"nonexistent\")\n\nlet test_store_latest_run () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" ~name:\"only\" () in\n Session.finish session ();\n let store = Store.open_ ~root () in\n let latest = Store.latest_run store () in\n is_true ~msg:\"latest returns the run\"\n (match latest with Some run -> Run.name run = Some \"only\" | None -> false)\n\nlet test_store_latest_run_with_filter () =\n with_temp_dir @@ fun root ->\n let s1 =\n Session.start ~root ~experiment:\"exp\" ~tags:[ \"train\" ] ~name:\"a\" ()\n in\n Session.finish s1 ();\n Unix.sleepf 0.01;\n let s2 =\n Session.start ~root ~experiment:\"exp\" ~tags:[ \"eval\" ] ~name:\"b\" ()\n in\n Session.finish s2 ();\n let store = Store.open_ ~root () in\n let latest = Store.latest_run store ~tag:\"train\" () in\n is_true ~msg:\"latest with tag filter\"\n (match latest with Some run -> Run.name run = Some \"a\" | None -> false)\n\nlet test_store_delete_run () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n let run = Session.run session in\n let store = Store.open_ ~root () in\n is_true ~msg:\"run dir exists before\" (Sys.file_exists (Run.dir run));\n Store.delete_run store run;\n is_false ~msg:\"run dir removed\" (Sys.file_exists (Run.dir run))\n\nlet test_store_delete_run_cleans_experiment () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"temp\" () in\n Session.finish session ();\n let run = Session.run session in\n let store = Store.open_ ~root () in\n Store.delete_run store run;\n let exp_dir = Filename.concat (Filename.concat root \"experiments\") \"temp\" in\n is_false ~msg:\"experiment dir removed\" (Sys.file_exists exp_dir)\n\nlet test_store_gc () =\n with_temp_dir @@ fun root ->\n (* Use directory artifact so blob is a directory (gc uses list_dirs) *)\n let dir = make_dir root \"dataset\" [ (\"a.txt\", \"data\") ] in\n let session = Session.start ~root ~experiment:\"exp\" () in\n let artifact =\n Session.log_artifact session ~name:\"data\" ~kind:`dataset ~path:dir ()\n in\n Session.finish session ();\n let store = Store.open_ ~root () in\n (* Blob is referenced, gc should remove 0 *)\n let removed_before = Store.gc store in\n equal ~msg:\"no unreferenced before\" int 0 removed_before;\n is_true ~msg:\"blob exists\" (Sys.file_exists (Artifact.path artifact));\n (* Remove artifact manifest to make blob unreferenced *)\n let version_dir =\n Filename.concat\n (Filename.concat\n (Filename.concat (Filename.concat root \"artifacts\") \"data\")\n \"versions\")\n \"v1\"\n in\n ignore (Sys.command (\"rm -rf \" ^ Filename.quote version_dir));\n let removed_after = Store.gc store in\n equal ~msg:\"one blob removed\" int 1 removed_after\n\nlet test_store_find_artifact () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"weights\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n ignore (Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src ());\n Session.finish session ();\n let store = Store.open_ ~root () in\n is_some ~msg:\"found\" (Store.find_artifact store ~name:\"model\" ~version:\"v1\")\n\nlet test_store_list_artifacts () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"model.bin\" \"weights\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n ignore (Session.log_artifact session ~name:\"model\" ~kind:`model ~path:src ());\n ignore (Session.log_artifact session ~name:\"data\" ~kind:`dataset ~path:src ());\n Session.finish session ();\n let store = Store.open_ ~root () in\n equal ~msg:\"all artifacts\" int 2 (List.length (Store.list_artifacts store ()));\n equal ~msg:\"filter by kind\" int 1\n (List.length (Store.list_artifacts store ~kind:`model ()));\n equal ~msg:\"filter by name\" int 1\n (List.length (Store.list_artifacts store ~name:\"data\" ()))\n\nlet store_tests =\n [\n test \"open creates dirs\" test_store_open;\n test \"list_experiments\" test_store_list_experiments;\n test \"list_runs all\" test_store_list_runs_all;\n test \"list_runs by experiment\" test_store_list_runs_experiment;\n test \"find_run\" test_store_find_run;\n test \"find_run missing\" test_store_find_run_missing;\n test \"latest_run\" test_store_latest_run;\n test \"latest_run with filter\" test_store_latest_run_with_filter;\n test \"delete_run\" test_store_delete_run;\n test \"delete_run cleans experiment\" test_store_delete_run_cleans_experiment;\n test \"gc\" test_store_gc;\n test \"find_artifact\" test_store_find_artifact;\n test \"list_artifacts\" test_store_list_artifacts;\n ]\n\n(* Run monitor *)\n\nlet test_monitor_poll () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"train/loss\" 1.0;\n Session.log_metric session ~step:2 \"train/loss\" 0.5;\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n let metrics = Run_monitor.metrics monitor in\n equal ~msg:\"one key\" int 1 (List.length metrics);\n let latest = List.assoc \"train/loss\" metrics in\n equal ~msg:\"latest step\" int 2 latest.step;\n Run_monitor.close monitor;\n Session.finish session ()\n\nlet test_monitor_incremental () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 1.0;\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n equal ~msg:\"one point after first poll\" int 1\n (List.length (Run_monitor.history monitor \"x\"));\n Session.log_metric session ~step:2 \"x\" 2.0;\n Run_monitor.poll monitor;\n equal ~msg:\"two points after second poll\" int 2\n (List.length (Run_monitor.history monitor \"x\"));\n Run_monitor.close monitor;\n Session.finish session ()\n\nlet test_monitor_history_chronological () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 3.0;\n Session.log_metric session ~step:2 \"x\" 1.0;\n Session.log_metric session ~step:3 \"x\" 2.0;\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n let history = Run_monitor.history monitor \"x\" in\n let steps = List.map fst history in\n equal ~msg:\"chronological steps\" (list int) [ 1; 2; 3 ] steps;\n Run_monitor.close monitor;\n Session.finish session ()\n\nlet test_monitor_metric_defs () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"loss\" ~summary:`Min ~goal:`Minimize ();\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n let defs = Run_monitor.metric_defs monitor in\n equal ~msg:\"one def\" int 1 (List.length defs);\n let def = List.assoc \"loss\" defs in\n is_true ~msg:\"summary min\" (def.summary = `Min);\n Run_monitor.close monitor;\n Session.finish session ()\n\nlet test_monitor_best_minimize () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"loss\" ~goal:`Minimize ();\n Session.log_metric session ~step:1 \"loss\" 1.0;\n Session.log_metric session ~step:2 \"loss\" 0.3;\n Session.log_metric session ~step:3 \"loss\" 0.7;\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n let best = Run_monitor.best monitor \"loss\" in\n is_true ~msg:\"best is 0.3\"\n (match best with Some m -> m.value = 0.3 | None -> false);\n Run_monitor.close monitor;\n Session.finish session ()\n\nlet test_monitor_best_maximize () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"acc\" ~goal:`Maximize ();\n Session.log_metric session ~step:1 \"acc\" 0.5;\n Session.log_metric session ~step:2 \"acc\" 0.9;\n Session.log_metric session ~step:3 \"acc\" 0.7;\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n let best = Run_monitor.best monitor \"acc\" in\n is_true ~msg:\"best is 0.9\"\n (match best with Some m -> m.value = 0.9 | None -> false);\n Run_monitor.close monitor;\n Session.finish session ()\n\nlet test_monitor_best_loss_heuristic () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"train/loss\" 1.0;\n Session.log_metric session ~step:2 \"train/loss\" 0.2;\n Session.log_metric session ~step:3 \"train/loss\" 0.5;\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n let best = Run_monitor.best monitor \"train/loss\" in\n is_true ~msg:\"loss heuristic picks minimum\"\n (match best with Some m -> m.value = 0.2 | None -> false);\n Run_monitor.close monitor;\n Session.finish session ()\n\nlet test_monitor_live_status () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 1.0;\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n is_true ~msg:\"live during session\" (Run_monitor.live_status monitor = `Live);\n Session.finish session ();\n Run_monitor.poll monitor;\n is_true ~msg:\"done after finish\"\n (match Run_monitor.live_status monitor with `Done _ -> true | _ -> false);\n Run_monitor.close monitor\n\nlet test_monitor_best_nonexistent () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n let run = Session.run session in\n let monitor = Run_monitor.start run in\n Run_monitor.poll monitor;\n is_none ~msg:\"no best for missing key\" (Run_monitor.best monitor \"nope\");\n Run_monitor.close monitor;\n Session.finish session ()\n\nlet run_monitor =\n [\n test \"poll reads metrics\" test_monitor_poll;\n test \"incremental polling\" test_monitor_incremental;\n test \"history chronological\" test_monitor_history_chronological;\n test \"metric_defs\" test_monitor_metric_defs;\n test \"best minimize\" test_monitor_best_minimize;\n test \"best maximize\" test_monitor_best_maximize;\n test \"best loss heuristic\" test_monitor_best_loss_heuristic;\n test \"live_status\" test_monitor_live_status;\n test \"best nonexistent key\" test_monitor_best_nonexistent;\n ]\n\n(* Robustness *)\n\nlet test_partial_log () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 1.0;\n let run = Session.run session in\n (* Append truncated JSON *)\n let oc =\n open_out_gen\n [ Open_append; Open_creat ]\n 0o644\n (Filename.concat (Run.dir run) \"events.jsonl\")\n in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc \"{\\\"type\\\":\\\"metric\\\"\");\n let run = Session.run session in\n equal ~msg:\"partial line ignored\" int 1\n (List.length (Run.metric_history run \"x\"));\n Session.finish session ()\n\nlet test_malformed_line () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 1.0;\n let run = Session.run session in\n let oc =\n open_out_gen\n [ Open_append; Open_creat ]\n 0o644\n (Filename.concat (Run.dir run) \"events.jsonl\")\n in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () -> output_string oc \"this is not json at all\\n\");\n let run = Session.run session in\n equal ~msg:\"malformed line skipped\" int 1\n (List.length (Run.metric_history run \"x\"));\n Session.finish session ()\n\nlet test_empty_event_log () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n let run = Session.run session in\n (* No events logged *)\n equal ~msg:\"no metrics\" int 0 (List.length (Run.metric_keys run));\n equal ~msg:\"no tags\" (list string) [] (Run.tags run);\n is_true ~msg:\"running\" (Run.status run = `running);\n Session.finish session ()\n\nlet test_mixed_valid_invalid () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_metric session ~step:1 \"x\" 1.0;\n let run = Session.run session in\n let oc =\n open_out_gen\n [ Open_append; Open_creat ]\n 0o644\n (Filename.concat (Run.dir run) \"events.jsonl\")\n in\n Fun.protect\n ~finally:(fun () -> close_out oc)\n (fun () ->\n output_string oc \"garbage\\n\";\n output_string oc \"{\\\"bad\\\":true}\\n\");\n Session.log_metric session ~step:2 \"x\" 2.0;\n let run = Session.run session in\n equal ~msg:\"only valid metrics\" int 2\n (List.length (Run.metric_history run \"x\"));\n Session.finish session ()\n\nlet robustness =\n [\n test \"partial log\" test_partial_log;\n test \"malformed line\" test_malformed_line;\n test \"empty event log\" test_empty_event_log;\n test \"mixed valid and invalid\" test_mixed_valid_invalid;\n ]\n\n(* Auto-computed summaries *)\n\nlet summary_float run key =\n match Run.find_summary run key with\n | Some v -> Value.to_float v\n | None -> None\n\nlet test_auto_summary_min () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"loss\" ~summary:`Min ();\n Session.log_metric session ~step:1 \"loss\" 1.0;\n Session.log_metric session ~step:2 \"loss\" 0.3;\n Session.log_metric session ~step:3 \"loss\" 0.7;\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"min summary\"\n (option (float 0.0))\n (Some 0.3) (summary_float run \"loss\")\n\nlet test_auto_summary_max () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"acc\" ~summary:`Max ();\n Session.log_metric session ~step:1 \"acc\" 0.5;\n Session.log_metric session ~step:2 \"acc\" 0.9;\n Session.log_metric session ~step:3 \"acc\" 0.7;\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"max summary\"\n (option (float 0.0))\n (Some 0.9) (summary_float run \"acc\")\n\nlet test_auto_summary_mean () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"x\" ~summary:`Mean ();\n Session.log_metric session ~step:1 \"x\" 1.0;\n Session.log_metric session ~step:2 \"x\" 2.0;\n Session.log_metric session ~step:3 \"x\" 3.0;\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"mean summary\"\n (option (float 0.0))\n (Some 2.0) (summary_float run \"x\")\n\nlet test_auto_summary_last () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"x\" ~summary:`Last ();\n Session.log_metric session ~step:1 \"x\" 1.0;\n Session.log_metric session ~step:2 \"x\" 2.0;\n Session.log_metric session ~step:3 \"x\" 3.0;\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"last summary\"\n (option (float 0.0))\n (Some 3.0) (summary_float run \"x\")\n\nlet test_auto_summary_none () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"x\" ~summary:`None ();\n Session.log_metric session ~step:1 \"x\" 1.0;\n Session.finish session ();\n let run = Session.run session in\n is_none ~msg:\"no auto-summary\" (Run.find_summary run \"x\")\n\nlet test_explicit_summary_wins () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.define_metric session \"loss\" ~summary:`Min ();\n Session.log_metric session ~step:1 \"loss\" 1.0;\n Session.log_metric session ~step:2 \"loss\" 0.3;\n Session.set_summary session [ (\"loss\", `Float 999.0) ];\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"explicit wins\"\n (option (float 0.0))\n (Some 999.0) (summary_float run \"loss\")\n\nlet auto_summaries =\n [\n test \"auto-summary min\" test_auto_summary_min;\n test \"auto-summary max\" test_auto_summary_max;\n test \"auto-summary mean\" test_auto_summary_mean;\n test \"auto-summary last\" test_auto_summary_last;\n test \"auto-summary none\" test_auto_summary_none;\n test \"explicit summary wins\" test_explicit_summary_wins;\n ]\n\n(* Grouping *)\n\nlet test_group_round_trip () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" ~group:\"sweep-lr\" () in\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"group\" (option string) (Some \"sweep-lr\") (Run.group run)\n\nlet test_group_none_by_default () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n equal ~msg:\"no group\" (option string) None (Run.group (Session.run session))\n\nlet test_run_list_filter_group () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp\" ~group:\"a\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp\" ~group:\"b\" () in\n Session.finish s2 ();\n let s3 = Session.start ~root ~experiment:\"exp\" () in\n Session.finish s3 ();\n equal ~msg:\"group a\" int 1\n (List.length (Run.list ~root ~experiment:\"exp\" ~group:\"a\" ()));\n equal ~msg:\"group b\" int 1\n (List.length (Run.list ~root ~experiment:\"exp\" ~group:\"b\" ()));\n equal ~msg:\"all\" int 3 (List.length (Run.list ~root ~experiment:\"exp\" ()))\n\nlet test_store_list_runs_group () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp\" ~group:\"sweep\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp\" () in\n Session.finish s2 ();\n let store = Store.open_ ~root () in\n equal ~msg:\"sweep only\" int 1\n (List.length (Store.list_runs store ~group:\"sweep\" ()))\n\nlet test_store_latest_run_group () =\n with_temp_dir @@ fun root ->\n let s1 = Session.start ~root ~experiment:\"exp\" ~group:\"a\" () in\n Session.finish s1 ();\n let s2 = Session.start ~root ~experiment:\"exp\" ~group:\"b\" () in\n Session.finish s2 ();\n let store = Store.open_ ~root () in\n let latest = Store.latest_run store ~group:\"a\" () in\n is_true ~msg:\"latest in group a\"\n (match latest with Some r -> Run.group r = Some \"a\" | None -> false)\n\nlet grouping =\n [\n test \"group round-trip\" test_group_round_trip;\n test \"group none by default\" test_group_none_by_default;\n test \"run list filter group\" test_run_list_filter_group;\n test \"store list_runs group\" test_store_list_runs_group;\n test \"store latest_run group\" test_store_latest_run_group;\n ]\n\n(* Media *)\n\nlet test_log_media_copies_file () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"pred.png\" \"image data\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_media session ~step:1 ~key:\"predictions\" ~kind:`Image ~path:src;\n Session.finish session ();\n let run = Session.run session in\n let entries = Run.media_history run \"predictions\" in\n equal ~msg:\"one entry\" int 1 (List.length entries);\n let entry = List.hd entries in\n is_true ~msg:\"file exists\" (Sys.file_exists entry.path);\n is_true ~msg:\"kind is image\" (entry.kind = `Image);\n equal ~msg:\"step\" int 1 entry.step\n\nlet test_log_media_nested_key () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"img.png\" \"pixels\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_media session ~step:5 ~key:\"train/predictions\" ~kind:`Image\n ~path:src;\n Session.finish session ();\n let run = Session.run session in\n let entries = Run.media_history run \"train/predictions\" in\n equal ~msg:\"one entry\" int 1 (List.length entries);\n let entry = List.hd entries in\n is_true ~msg:\"file in subdir\"\n (let parts = String.split_on_char '/' entry.path in\n List.exists (String.equal \"train\") parts)\n\nlet test_log_media_multiple_steps () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"img.png\" \"pixels\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_media session ~step:1 ~key:\"pred\" ~kind:`Image ~path:src;\n Session.log_media session ~step:5 ~key:\"pred\" ~kind:`Image ~path:src;\n Session.log_media session ~step:10 ~key:\"pred\" ~kind:`Image ~path:src;\n Session.finish session ();\n let run = Session.run session in\n let entries = Run.media_history run \"pred\" in\n equal ~msg:\"three entries\" int 3 (List.length entries);\n let steps = List.map (fun (e : Run.media_entry) -> e.step) entries in\n equal ~msg:\"chronological\" (list int) [ 1; 5; 10 ] steps\n\nlet test_log_media_missing_path_raises () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n raises_invalid_arg\n \"Munin.Session.log_media: path does not exist: /no/such/file\" (fun () ->\n Session.log_media session ~step:1 ~key:\"x\" ~kind:`File\n ~path:\"/no/such/file\");\n Session.finish session ()\n\nlet test_log_media_closed_ignored () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"img.png\" \"pixels\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n Session.log_media session ~step:1 ~key:\"pred\" ~kind:`Image ~path:src;\n let run = Session.run session in\n equal ~msg:\"no media\" int 0 (List.length (Run.media_keys run))\n\nlet test_log_table () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_table session ~step:1 ~key:\"confusion\" ~columns:[ \"cat\"; \"dog\" ]\n ~rows:[ [ `Int 90; `Int 10 ]; [ `Int 5; `Int 95 ] ];\n Session.finish session ();\n let run = Session.run session in\n let entries = Run.media_history run \"confusion\" in\n equal ~msg:\"one entry\" int 1 (List.length entries);\n let entry = List.hd entries in\n is_true ~msg:\"kind is table\" (entry.kind = `Table);\n is_true ~msg:\"json file exists\" (Sys.file_exists entry.path)\n\nlet test_media_keys_sorted () =\n with_temp_dir @@ fun root ->\n let src = make_file root \"f.bin\" \"data\" in\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.log_media session ~step:1 ~key:\"z/output\" ~kind:`File ~path:src;\n Session.log_media session ~step:1 ~key:\"a/input\" ~kind:`File ~path:src;\n Session.log_media session ~step:1 ~key:\"m/middle\" ~kind:`File ~path:src;\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"sorted\" (list string)\n [ \"a/input\"; \"m/middle\"; \"z/output\" ]\n (Run.media_keys run)\n\nlet test_media_empty_history () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n Session.finish session ();\n let run = Session.run session in\n equal ~msg:\"no media keys\" (list string) [] (Run.media_keys run);\n equal ~msg:\"empty history\" int 0\n (List.length (Run.media_history run \"nonexistent\"))\n\nlet media =\n [\n test \"log_media copies file\" test_log_media_copies_file;\n test \"log_media nested key\" test_log_media_nested_key;\n test \"log_media multiple steps\" test_log_media_multiple_steps;\n test \"log_media missing path raises\" test_log_media_missing_path_raises;\n test \"log_media closed ignored\" test_log_media_closed_ignored;\n test \"log_table\" test_log_table;\n test \"media_keys sorted\" test_media_keys_sorted;\n test \"media empty history\" test_media_empty_history;\n ]\n\n(* System monitor *)\n\nlet test_system_monitor_logs_metrics () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n let monitor = Munin_sys.start session ~interval:0.1 () in\n Thread.delay 0.35;\n Munin_sys.stop monitor;\n Session.finish session ();\n let run = Session.run session in\n let keys = Run.metric_keys run in\n is_true ~msg:\"has sys/cpu_user\" (List.mem \"sys/cpu_user\" keys);\n is_true ~msg:\"has sys/mem_used_pct\" (List.mem \"sys/mem_used_pct\" keys);\n is_true ~msg:\"has sys/proc_mem_mb\" (List.mem \"sys/proc_mem_mb\" keys)\n\nlet test_system_monitor_defines_metrics () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n let monitor = Munin_sys.start session ~interval:100.0 () in\n Munin_sys.stop monitor;\n Session.finish session ();\n let run = Session.run session in\n let defs = Run.metric_defs run in\n let has_def key =\n match List.assoc_opt key defs with\n | Some d -> d.summary = `Last\n | None -> false\n in\n is_true ~msg:\"cpu_user def\" (has_def \"sys/cpu_user\");\n is_true ~msg:\"mem_used_pct def\" (has_def \"sys/mem_used_pct\");\n is_true ~msg:\"proc_mem_mb def\" (has_def \"sys/proc_mem_mb\")\n\nlet test_system_monitor_stop_idempotent () =\n with_temp_dir @@ fun root ->\n let session = Session.start ~root ~experiment:\"exp\" () in\n let monitor = Munin_sys.start session ~interval:100.0 () in\n Munin_sys.stop monitor;\n Munin_sys.stop monitor;\n Session.finish session ()\n\nlet system_monitor_tests =\n [\n test \"logs metrics\" test_system_monitor_logs_metrics;\n test \"defines metrics\" test_system_monitor_defines_metrics;\n test \"stop idempotent\" test_system_monitor_stop_idempotent;\n ]\n\n(* Suite *)\n\nlet suite =\n [\n group \"Session lifecycle\" lifecycle;\n group \"Scalars\" scalars;\n group \"Metric definitions\" metric_definitions;\n group \"Metadata\" metadata;\n group \"Provenance\" provenance_tests;\n group \"Run loading\" run_loading;\n group \"Artifacts\" artifacts;\n group \"Store\" store_tests;\n group \"Run monitor\" run_monitor;\n group \"Robustness\" robustness;\n group \"Auto-computed summaries\" auto_summaries;\n group \"Grouping\" grouping;\n group \"Media\" media;\n group \"System monitor\" system_monitor_tests;\n ]\n\nlet () = run \"Munin\" suite\n"
},
{
"path": "packages/norn/README.md",
"content": "# Norn\n\nMCMC sampling with automatic gradients for OCaml, powered by [Rune](../rune/)\n\nNorn provides Hamiltonian Monte Carlo and NUTS samplers that leverage Rune's\nautomatic differentiation. You supply an unnormalized log-density function and\nan initial position; Norn handles gradient computation, trajectory integration,\nand Stan-style window adaptation of step size and mass matrix. Common workflows\nare one-line calls, while the kernel API gives full control over the sampling\npipeline.\n\n## Quick Start\n\nSample from a 2D Gaussian with NUTS:\n\n```ocaml\nopen Nx\n\nlet () =\n Rng.run ~seed:42 @@ fun () ->\n let f64 = Nx.float64 in\n\n (* Target: N([3; -1], [[1, 0.8]; [0.8, 1]]) *)\n let mu = Nx.create f64 [| 2 |] [| 3.0; -1.0 |] in\n let prec = Nx.create f64 [| 2; 2 |] [| 5.0; -4.0; -4.0; 5.0 |] in\n let log_prob x =\n let d = Nx.sub x mu in\n let d_col = Nx.reshape [| 2; 1 |] d in\n Nx.mul_s (Nx.squeeze (Nx.matmul (Nx.matrix_transpose d_col) (Nx.matmul prec d_col))) (-0.5)\n in\n\n let init = Nx.zeros f64 [| 2 |] in\n let result = Norn.nuts ~n:1000 log_prob init in\n\n let mean = Nx.mean ~axes:[ 0 ] result.samples in\n Printf.printf \"posterior mean: %s\\n\" (Nx.data_to_string mean);\n Printf.printf \"accept rate: %.2f\\n\" result.stats.accept_rate;\n Printf.printf \"ESS: %s\\n\" (Nx.data_to_string (Norn.ess result.samples))\n```\n\n## Features\n\n- **One-line sampling**: `Norn.hmc` and `Norn.nuts` for common workflows\n- **Configurable API**: `Norn.sample` with custom kernels via `make_kernel`\n- **Automatic gradients**: log-density gradients computed by Rune -- no manual derivatives\n- **Symplectic integrators**: `leapfrog`, `mclachlan`, `yoshida`\n- **Mass matrix metrics**: `unit_metric`, `diagonal_metric`, `dense_metric`\n- **Stan-style adaptation**: dual averaging for step size, Welford estimation for mass matrix\n- **Diagnostics**: effective sample size (`ess`) and split R-hat (`rhat`)\n\n## Examples\n\n- **01-sampling-basics** -- Sample from a correlated 2D Gaussian with NUTS\n- **02-bayesian-regression** -- Bayesian linear regression with credible intervals\n- **03-diagnostics** -- Multi-chain convergence checking with ESS and R-hat\n\n## Contributing\n\nSee the [Raven monorepo README](../README.md) for guidelines.\n\n## License\n\nISC License. See [LICENSE](../LICENSE) for details.\n"
},
{
"path": "packages/norn/bench/bench_norn.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet f64 = Nx.float64\n\nlet make_log_prob dim =\n let mean = Nx.zeros f64 [| dim |] in\n fun x ->\n let d = Nx.sub x mean in\n Nx.mul_s (Nx.sum (Nx.mul d d)) (-0.5)\n\nlet hmc_benches () =\n let cases = [ (\"2D_3lf\", 2, 3); (\"5D_3lf\", 5, 3) ] in\n List.map\n (fun (label, dim, num_leapfrog) ->\n let log_prob = make_log_prob dim in\n let metric = Norn.unit_metric dim in\n let kernel = Norn.hmc_kernel ~num_leapfrog ~step_size:0.1 ~metric () in\n let init = Nx.zeros f64 [| dim |] in\n let state = ref (kernel.init init log_prob) in\n Thumper.bench (Printf.sprintf \"HMC/%s\" label) (fun () ->\n let new_state, info = kernel.step !state log_prob in\n state := new_state;\n info.Norn.acceptance_rate))\n cases\n\nlet nuts_benches () =\n let cases = [ (\"2D_d3\", 2, 3); (\"5D_d3\", 5, 3) ] in\n List.map\n (fun (label, dim, max_depth) ->\n let log_prob = make_log_prob dim in\n let metric = Norn.unit_metric dim in\n let kernel = Norn.nuts_kernel ~max_depth ~step_size:0.1 ~metric () in\n let init = Nx.zeros f64 [| dim |] in\n let state = ref (kernel.init init log_prob) in\n Thumper.bench (Printf.sprintf \"NUTS/%s\" label) (fun () ->\n let new_state, info = kernel.step !state log_prob in\n state := new_state;\n info.Norn.acceptance_rate))\n cases\n\nlet ess_benches () =\n let cases = [ (\"2D_n100\", 2, 100); (\"5D_n100\", 5, 100) ] in\n List.map\n (fun (label, dim, n) ->\n let samples = Nx.randn f64 [| n; dim |] in\n Thumper.bench (Printf.sprintf \"ESS/%s\" label) (fun () -> Norn.ess samples))\n cases\n\nlet rhat_benches () =\n let cases = [ (\"2D_n100\", 2, 100); (\"5D_n100\", 5, 100) ] in\n List.map\n (fun (label, dim, n) ->\n let chains = Array.init 4 (fun _ -> Nx.randn f64 [| n; dim |]) in\n Thumper.bench (Printf.sprintf \"Rhat/%s\" label) (fun () ->\n Norn.rhat chains))\n cases\n\nlet () =\n Nx.Rng.run ~seed:42 (fun () ->\n Thumper.run \"norn\"\n [\n Thumper.group \"HMC\" (hmc_benches ());\n Thumper.group \"NUTS\" (nuts_benches ());\n Thumper.group \"ESS\" (ess_benches ());\n Thumper.group \"Rhat\" (rhat_benches ());\n ])\n"
},
{
"path": "packages/norn/bench/dune",
"content": "(executable\n (name bench_norn)\n (libraries nx rune norn thumper))\n\n(rule\n (alias runtest)\n (action\n (progn\n (run %{exe:bench_norn.exe} -q)\n (diff? norn.thumper norn.thumper.corrected))))\n"
},
{
"path": "packages/norn/bench/norn.thumper",
"content": "# thumper baseline\n# version: 1\n# suite_name: norn\n# host: 1480401c3b76ed18\n# cpu: Apple M1 Max\n# ocaml: 5.4.1\n# git: 31747323\n# dirty: true\n# command: /Users/tmattio/Workspace/raven/_build/default/packages/norn/bench/bench_norn.exe --bless --quick\n\ness/ess_2d_n100\talloc_words\t7.352397e+06\t7.352397e+06\t7.352397e+06\t0.000000e+00\t5\t1\ness/ess_2d_n100\tcpu_time\t1.559885e-02\t1.557237e-02\t1.567803e-02\t3.386704e-03\t5\t0\ness/ess_2d_n100\twall_time\t1.560669e-02\t1.557799e-02\t1.567933e-02\t3.246829e-03\t5\t0\ness/ess_5d_n100\talloc_words\t1.837855e+07\t1.837855e+07\t1.837855e+07\t0.000000e+00\t5\t1\ness/ess_5d_n100\tcpu_time\t4.013211e-02\t3.996071e-02\t4.034314e-02\t4.764651e-03\t5\t0\ness/ess_5d_n100\twall_time\t4.025803e-02\t4.005148e-02\t4.049841e-02\t5.550847e-03\t5\t0\nhmc/hmc_2d_3lf\talloc_words\t2.538700e+04\t2.538700e+04\t2.538700e+04\t0.000000e+00\t5\t0\nhmc/hmc_2d_3lf\tcpu_time\t1.123211e-04\t1.114248e-04\t1.134997e-04\t9.236276e-03\t5\t1\nhmc/hmc_2d_3lf\twall_time\t1.124071e-04\t1.114420e-04\t1.136754e-04\t9.934176e-03\t5\t1\nhmc/hmc_5d_3lf\talloc_words\t2.548300e+04\t2.548200e+04\t2.548300e+04\t1.962092e-05\t5\t0\nhmc/hmc_5d_3lf\tcpu_time\t1.165771e-04\t1.158171e-04\t1.172042e-04\t5.949159e-03\t5\t0\nhmc/hmc_5d_3lf\twall_time\t1.167592e-04\t1.160356e-04\t1.174092e-04\t5.882306e-03\t5\t0\nnuts/nuts_2d_d3\talloc_words\t8.002300e+04\t5.842900e+04\t8.409700e+04\t1.603789e-01\t5\t0\nnuts/nuts_2d_d3\tcpu_time\t3.536361e-04\t3.497111e-04\t3.602631e-04\t1.491926e-02\t5\t0\nnuts/nuts_2d_d3\twall_time\t3.537679e-04\t3.499804e-04\t3.604565e-04\t1.480642e-02\t5\t0\nnuts/nuts_5d_d3\talloc_words\t8.291500e+04\t8.290400e+04\t8.291900e+04\t9.045408e-05\t5\t1\nnuts/nuts_5d_d3\tcpu_time\t3.620090e-04\t3.591001e-04\t3.636080e-04\t6.226318e-03\t5\t0\nnuts/nuts_5d_d3\twall_time\t3.621874e-04\t3.585911e-04\t3.638656e-04\t7.281394e-03\t5\t0\nrhat/rhat_2d_n100\talloc_words\t3.096500e+04\t3.096500e+04\t3.096500e+04\t0.000000e+00\t5\t0\nrhat/rhat_2d_n100\tcpu_time\t1.348335e-04\t1.338291e-04\t1.359049e-04\t7.697749e-03\t5\t1\nrhat/rhat_2d_n100\twall_time\t1.351109e-04\t1.339526e-04\t1.364586e-04\t9.274007e-03\t5\t1\nrhat/rhat_5d_n100\talloc_words\t3.096500e+04\t3.096500e+04\t3.096500e+04\t0.000000e+00\t5\t0\nrhat/rhat_5d_n100\tcpu_time\t1.453519e-04\t1.445829e-04\t1.461137e-04\t5.265860e-03\t5\t0\nrhat/rhat_5d_n100\twall_time\t1.457385e-04\t1.450074e-04\t1.464628e-04\t4.993161e-03\t5\t0\n"
},
{
"path": "packages/norn/doc/01-getting-started.md",
"content": "# Getting Started\n\nThis guide shows you how to sample from a target distribution using Norn's\nMCMC samplers.\n\n## Installation\n\n\n```bash\nopam install norn\n```\n\nOr build from source:\n\n\n```bash\ngit clone https://github.com/raven-ml/raven\ncd raven && dune build norn\n```\n\nAdd to your `dune` file:\n\n\n```dune\n(executable\n (name main)\n (libraries norn nx rune))\n```\n\n## Your First Sampler\n\nNorn samplers take three things: a sample count, an unnormalized log-density\nfunction, and an initial position. Here we sample from a 2D standard Gaussian\nusing NUTS:\n\n\n```ocaml\nopen Nx\n\nlet () =\n Rng.run ~seed:42 @@ fun () ->\n let f64 = Nx.float64 in\n\n (* log p(x) = -0.5 * ||x||^2 (standard Gaussian) *)\n let log_prob x = Nx.mul_s (Nx.sum (Nx.square x)) (-0.5) in\n\n let init = Nx.zeros f64 [| 2 |] in\n let result = Norn.nuts ~n:1000 log_prob init in\n\n Printf.printf \"samples shape: %s\\n\"\n (String.concat \"x\" (List.map string_of_int\n (Array.to_list (Nx.shape result.samples))));\n Printf.printf \"accept rate: %.3f\\n\" result.stats.accept_rate;\n Printf.printf \"divergences: %d\\n\" result.stats.num_divergent\n```\n\nKey points:\n- `log_prob` returns a scalar `Nx.float64_t` (not a float) -- Rune differentiates it automatically\n- `init` is the starting position, shape `[dim]`\n- `result.samples` has shape `[n; dim]` -- one row per sample\n- NUTS adapts trajectory length automatically via U-turn detection\n\n## Understanding the Result\n\n`Norn.nuts` and `Norn.hmc` return a `result` record:\n\n\n```ocaml\ntype result = {\n samples : Nx.float64_t; (* shape [n; dim] *)\n log_densities : Nx.float64_t; (* shape [n] *)\n stats : stats;\n}\n\ntype stats = {\n accept_rate : float; (* mean acceptance rate during sampling *)\n step_size : float; (* final adapted step size *)\n num_divergent : int; (* number of divergent transitions *)\n}\n```\n\nCompute posterior summaries from `result.samples`:\n\n\n```ocaml\nlet mean = Nx.mean ~axes:[ 0 ] result.samples in\nlet std = Nx.std ~axes:[ 0 ] result.samples in\nPrintf.printf \"mean: %s\\n\" (Nx.data_to_string mean);\nPrintf.printf \"std: %s\\n\" (Nx.data_to_string std)\n```\n\n## Using HMC\n\nHMC requires a fixed number of leapfrog steps per transition. It is simpler\nthan NUTS but requires tuning `num_leapfrog`:\n\n\n```ocaml\nlet result =\n Norn.hmc ~n:1000 ~num_leapfrog:30 log_prob init\n```\n\nDefault values: `step_size = 0.01`, `target_accept = 0.65`, `num_leapfrog = 20`.\nStep size and mass matrix are adapted during warmup regardless of the sampler.\n\n## The Kernel API\n\nFor more control, use `Norn.sample` with a kernel constructor. The `make_kernel`\nfunction receives adapted step size and metric at each warmup step:\n\n\n```ocaml\nlet result =\n Norn.sample ~n:1000 log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~step_size ~metric ())\n```\n\nThis is equivalent to `Norn.nuts ~n:1000 log_prob init`, but you can customize\nthe kernel:\n\n\n```ocaml\nlet result =\n Norn.sample ~n:1000 log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~integrator:Norn.mclachlan ~max_depth:8\n ~step_size ~metric ())\n```\n\nThe `make_kernel` signature is `step_size:float -> metric:metric -> kernel`.\nDuring warmup, `sample` calls `make_kernel` each step with the latest adapted\nvalues. After warmup, it freezes the final step size and metric for all\nsampling steps.\n\n## HMC vs NUTS\n\n| Aspect | HMC | NUTS |\n|--------|-----|------|\n| Trajectory length | Fixed (`num_leapfrog` steps) | Automatic (U-turn detection) |\n| Tuning parameters | `step_size`, `num_leapfrog` | `step_size`, `max_depth` |\n| Default target accept | 0.65 | 0.80 |\n| Gradient evaluations | `num_leapfrog` per step | Variable, up to `2^max_depth` |\n| Best for | Simple, well-conditioned posteriors | General use |\n\nNUTS is the recommended default. Use HMC when you know the optimal trajectory\nlength or need predictable cost per step.\n\n## Next Steps\n\n- [Adaptation and Diagnostics](../02-adaptation-and-diagnostics/) -- warmup windows, ESS, R-hat\n- [Advanced Usage](../03-advanced-usage/) -- custom integrators, metrics, and monitoring\n- [PyMC Comparison](../04-pymc-comparison/) -- mapping from Python's PyMC/BlackJAX to Norn\n"
},
{
"path": "packages/norn/doc/02-adaptation-and-diagnostics.md",
"content": "# Adaptation and Diagnostics\n\nNorn uses Stan-style window adaptation during warmup to tune step size and\nmass matrix automatically. After sampling, diagnostics like ESS and R-hat\nhelp you assess whether the chain has converged.\n\n## Window Adaptation\n\nWhen you call `Norn.nuts`, `Norn.hmc`, or `Norn.sample`, the first\n`num_warmup` iterations (default `n / 2`) are discarded as warmup. During\nwarmup, Norn adapts two quantities:\n\n1. **Step size** -- via dual averaging (Nesterov 2009)\n2. **Mass matrix** -- via regularized Welford variance estimation\n\nWarmup is divided into three phases following Stan's scheme:\n\n| Phase | What adapts | Description |\n|-------|-------------|-------------|\n| Initial fast | Step size only | Short burn-in to find a reasonable step size |\n| Slow windows | Step size + mass matrix | Doubling windows that collect samples for Welford estimation. At the end of each window, the mass matrix is updated and step size is reset |\n| Final fast | Step size only | Short phase to re-tune step size for the final metric |\n\nThe slow windows double in length (e.g., 25, 50, 100, ...) so more samples\ncontribute to later mass matrix estimates, which are more reliable as the\nchain moves closer to the typical set.\n\n## Step Size Adaptation\n\nStep size is tuned via dual averaging to reach a target acceptance rate.\nThe default targets are:\n\n- HMC: `target_accept = 0.65`\n- NUTS: `target_accept = 0.80`\n\nYou can override the target:\n\n\n```ocaml\nlet result = Norn.nuts ~n:1000 ~target_accept:0.90 log_prob init\n```\n\nHigher target acceptance rates produce smaller step sizes, which give more\naccurate trajectories at the cost of more computation. Values between 0.6\nand 0.9 work well for most problems.\n\nYou can also set the initial step size:\n\n\n```ocaml\nlet result = Norn.nuts ~n:1000 ~step_size:0.1 log_prob init\n```\n\nThe final adapted step size is available in `result.stats.step_size`.\n\n## Mass Matrix Adaptation\n\nThe mass matrix (inverse metric) controls the shape of the momentum\ndistribution. A well-chosen mass matrix makes the sampler's kinetic energy\nmatch the target's geometry, improving mixing.\n\nDuring slow windows, Norn collects position samples and estimates the\ninverse mass matrix as a diagonal covariance using the Welford online\nalgorithm with shrinkage regularization. At the end of each slow window:\n\n1. The diagonal inverse mass matrix is computed from accumulated statistics\n2. A `diagonal_metric` is constructed from the estimate\n3. The Welford accumulator is reset for the next window\n4. Step size dual averaging is reset to re-tune for the new metric\n\nAfter warmup, the metric is frozen and used for all sampling iterations.\n\n## Controlling Warmup\n\nSet `num_warmup` explicitly when the default (`n / 2`) is too few or too many:\n\n\n```ocaml\n(* More warmup for a difficult posterior *)\nlet result = Norn.nuts ~n:1000 ~num_warmup:2000 log_prob init\n\n(* Less warmup when the posterior is simple *)\nlet result = Norn.nuts ~n:1000 ~num_warmup:100 log_prob init\n```\n\n## Effective Sample Size\n\nAutocorrelated MCMC samples contain less information than independent samples.\nThe effective sample size (ESS) estimates how many independent samples the\nchain is worth:\n\n\n```ocaml\nlet result = Norn.nuts ~n:2000 log_prob init in\nlet n_eff = Norn.ess result.samples in\nPrintf.printf \"ESS: %s\\n\" (Nx.data_to_string n_eff)\n```\n\n`ess` takes a matrix of shape `[n; dim]` and returns a vector of shape `[dim]`\nwith the ESS for each parameter. It uses autocorrelation with the initial\nmonotone sequence estimator.\n\nRules of thumb:\n- ESS > 100 per parameter is often sufficient for posterior means\n- ESS > 400 is preferred for tail quantiles\n- Low ESS relative to `n` suggests poor mixing -- consider reparameterization\n or a different metric\n\n## Split R-hat\n\nR-hat measures convergence by comparing within-chain and between-chain variance.\nIt requires multiple chains:\n\n\n```ocaml\nopen Nx\n\nlet () =\n Rng.run ~seed:42 @@ fun () ->\n let f64 = Nx.float64 in\n let log_prob x = Nx.mul_s (Nx.sum (Nx.square x)) (-0.5) in\n\n (* Run 4 chains from different starting points *)\n let chains =\n Array.init 4 (fun i ->\n let init = Nx.mul_s (Nx.ones f64 [| 3 |]) (Float.of_int (i - 2)) in\n let result = Norn.nuts ~n:500 log_prob init in\n result.samples)\n in\n\n let r = Norn.rhat chains in\n Printf.printf \"R-hat: %s\\n\" (Nx.data_to_string r)\n```\n\n`rhat` takes an array of chains, each of shape `[n; dim]`, and returns\nshape `[dim]`. It uses the split R-hat variant (each chain is split in half\nbefore comparison).\n\nInterpretation:\n- R-hat close to 1.0 indicates convergence\n- R-hat > 1.01 suggests the chains have not mixed\n- R-hat > 1.1 is a strong signal of non-convergence\n\n## Checking Convergence\n\nA practical convergence check combines multiple diagnostics:\n\n\n```ocaml\nlet check_convergence (results : Norn.result array) =\n let chains = Array.map (fun r -> r.Norn.samples) results in\n let r = Norn.rhat chains in\n\n (* Check R-hat for all parameters *)\n let max_rhat = Nx.item [] (Nx.max r) in\n if max_rhat > 1.01 then\n Printf.printf \"WARNING: max R-hat = %.3f (chains have not converged)\\n\"\n max_rhat;\n\n (* Check ESS for each chain *)\n Array.iteri\n (fun i result ->\n let n_eff = Norn.ess result.Norn.samples in\n let min_ess = Nx.item [] (Nx.min n_eff) in\n Printf.printf \"chain %d: min ESS = %.0f, divergences = %d\\n\" i min_ess\n result.stats.num_divergent)\n results;\n\n (* Check divergences *)\n let total_div =\n Array.fold_left (fun acc r -> acc + r.Norn.stats.num_divergent) 0 results\n in\n if total_div > 0 then\n Printf.printf \"WARNING: %d divergent transitions (consider reparameterization)\\n\"\n total_div\n```\n\n## Next Steps\n\n- [Advanced Usage](../03-advanced-usage/) -- custom integrators, metrics, and monitoring\n- [Getting Started](../01-getting-started/) -- basic usage and the kernel API\n- [PyMC Comparison](../04-pymc-comparison/) -- mapping from Python's PyMC/BlackJAX to Norn\n"
},
{
"path": "packages/norn/doc/03-advanced-usage.md",
"content": "# Advanced Usage\n\nThis guide covers custom integrators, metrics, kernel composition via\n`Norn.sample`, and monitoring sampling progress.\n\n## Integrators\n\nThe integrator controls how Hamiltonian dynamics are approximated. Norn\nprovides three symplectic integrators:\n\n| Integrator | Order | Grad evals/step | Best for |\n|-----------|-------|-----------------|----------|\n| `leapfrog` | 2nd | 1 | General use (default) |\n| `mclachlan` | 2nd | 2 | Higher acceptance on stiff problems |\n| `yoshida` | 4th | 3 | High accuracy with fewer steps |\n\n### Leapfrog (default)\n\nThe standard velocity Verlet integrator. One gradient evaluation per step,\ngood balance of accuracy and cost:\n\n\n```ocaml\nlet result =\n Norn.sample ~n:1000 log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~integrator:Norn.leapfrog ~step_size ~metric ())\n```\n\n### McLachlan\n\nMcLachlan's two-stage integrator achieves higher acceptance rates than\nleapfrog on challenging posteriors at the cost of two gradient evaluations\nper step:\n\n\n```ocaml\nlet result =\n Norn.sample ~n:1000 log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~integrator:Norn.mclachlan ~step_size ~metric ())\n```\n\nUse McLachlan when leapfrog produces too many divergences or low acceptance\nrates despite adaptation.\n\n### Yoshida\n\nYoshida's fourth-order integrator is more accurate than leapfrog, allowing\nlarger step sizes or fewer integration steps. Three gradient evaluations\nper step:\n\n\n```ocaml\nlet result =\n Norn.sample ~n:1000 log_prob init (fun ~step_size ~metric ->\n Norn.hmc_kernel ~integrator:Norn.yoshida ~num_leapfrog:10\n ~step_size ~metric ())\n```\n\nYoshida is most useful with HMC where the trajectory length is fixed -- the\nhigher accuracy lets you use fewer steps for the same trajectory quality.\n\n## Metrics\n\nThe metric defines the mass matrix, which shapes the momentum distribution\nto match the target geometry. A good metric improves mixing by making the\nsampler's kinetic energy reflect the posterior's covariance structure.\n\n### Unit Metric\n\nIdentity mass matrix. Momentum sampled from `N(0, I)`. This is the starting\npoint for adaptation:\n\n\n```ocaml\nlet m = Norn.unit_metric dim\n```\n\n### Diagonal Metric\n\nDiagonal mass matrix estimated from the inverse variance of each parameter.\nThis is what window adaptation produces automatically:\n\n\n```ocaml\nlet f64 = Nx.float64 in\nlet inv_mass_diag = Nx.create f64 [| 2 |] [| 1.0; 0.01 |] in\nlet m = Norn.diagonal_metric inv_mass_diag\n```\n\nUse a diagonal metric when parameters have very different scales. Adaptation\nestimates this automatically, but you can provide your own if you know the\nposterior variances.\n\n### Dense Metric\n\nFull inverse mass matrix. Uses Cholesky decomposition for momentum sampling.\nCaptures correlations between parameters:\n\n\n```ocaml\nlet f64 = Nx.float64 in\nlet inv_mass =\n Nx.create f64 [| 2; 2 |] [| 1.0; 0.8; 0.8; 1.0 |]\nin\nlet m = Norn.dense_metric inv_mass\n```\n\nDense metrics help with strongly correlated posteriors but are expensive for\nhigh-dimensional problems (`O(dim^2)` storage, `O(dim^3)` Cholesky).\n\n## Composing Kernels with sample\n\n`Norn.sample` is the configurable entry point. The `make_kernel` function\nreceives the current adapted step size and metric, returning a kernel:\n\n\n```ocaml\nlet result =\n Norn.sample ~n:2000 ~num_warmup:1000 ~target_accept:0.85\n log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~integrator:Norn.mclachlan ~max_depth:8\n ~step_size ~metric ())\n```\n\nThis gives you full control over:\n- The sampler algorithm (HMC vs NUTS)\n- The integrator (leapfrog, mclachlan, yoshida)\n- Algorithm-specific parameters (`num_leapfrog`, `max_depth`)\n- Step size and metric are provided by adaptation\n\nThe `make_kernel` function is called at every warmup step (with updated\nadaptation values) and once more with the final values before sampling begins.\n\n## Monitoring with report\n\nThe `~report` callback lets you monitor sampling progress. It is called after\neach step with the current step number, state, and diagnostics:\n\n\n```ocaml\nlet report ~step state info =\n if step mod 100 = 0 then\n Printf.printf \"step %4d log_p = %.2f accept = %.3f steps = %d%s\\n\"\n step state.Norn.log_density info.Norn.acceptance_rate\n info.num_integration_steps\n (if info.is_divergent then \" DIVERGENT\" else \"\")\n\nlet result =\n Norn.sample ~n:1000 ~report log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~step_size ~metric ())\n```\n\nStep numbers are negative during warmup (counting down to zero) and\nnon-negative during sampling. This makes it easy to distinguish the two\nphases:\n\n\n```ocaml\nlet report ~step _state info =\n if step < 0 then\n Printf.printf \"warmup %4d accept = %.3f\\n\" step info.Norn.acceptance_rate\n else if step mod 100 = 0 then\n Printf.printf \"sample %4d accept = %.3f\\n\" step info.acceptance_rate\n```\n\n## Providing a Known Metric\n\nIf you know the posterior covariance from a previous run or analytic\ncalculation, skip the adaptation overhead by providing the metric directly:\n\n\n```ocaml\nopen Nx\n\nlet () =\n Rng.run ~seed:42 @@ fun () ->\n let f64 = Nx.float64 in\n let log_prob x = Nx.mul_s (Nx.sum (Nx.square x)) (-0.5) in\n let init = Nx.zeros f64 [| 2 |] in\n\n (* Use a known diagonal inverse mass *)\n let inv_mass_diag = Nx.create f64 [| 2 |] [| 1.0; 1.0 |] in\n let metric = Norn.diagonal_metric inv_mass_diag in\n\n let result =\n Norn.sample ~n:1000 ~num_warmup:200 log_prob init (fun ~step_size ~metric:_ ->\n Norn.nuts_kernel ~step_size ~metric ())\n in\n Printf.printf \"accept rate: %.3f\\n\" result.stats.accept_rate\n```\n\nNote that `~metric:_` ignores the adapted metric and uses the fixed one.\nStep size is still adapted during warmup.\n\n## Next Steps\n\n- [Getting Started](../01-getting-started/) -- basic usage and the kernel API\n- [Adaptation and Diagnostics](../02-adaptation-and-diagnostics/) -- warmup windows, ESS, R-hat\n- [PyMC Comparison](../04-pymc-comparison/) -- mapping from Python's PyMC/BlackJAX to Norn\n"
},
{
"path": "packages/norn/doc/04-pymc-comparison.md",
"content": "# PyMC Comparison\n\nThis page maps [PyMC](https://www.pymc.io/) and\n[BlackJAX](https://github.com/blackjax-devs/blackjax) concepts to their\nNorn equivalents. Norn's design is closest to BlackJAX: both provide\nfunctional kernel APIs where the sampler state is explicit and the log-density\nfunction is passed at each step.\n\n## One-Line Sampling\n\n**PyMC:**\n\n```python\nimport pymc as pm\n\nwith pm.Model():\n x = pm.Normal(\"x\", mu=0, sigma=1, shape=2)\n trace = pm.sample(1000, tune=500)\n```\n\n**Norn:**\n\n\n```ocaml\nlet log_prob x = Nx.mul_s (Nx.sum (Nx.square x)) (-0.5) in\nlet init = Nx.zeros Nx.float64 [| 2 |] in\nlet result = Norn.nuts ~n:1000 ~num_warmup:500 log_prob init\n```\n\nPyMC builds a probabilistic model and derives the log-density automatically.\nNorn takes the log-density function directly -- you write it yourself or\nbuild it from your model. Rune handles the gradient.\n\n## BlackJAX Kernel API\n\n**BlackJAX:**\n\n```python\nimport blackjax\nimport jax\n\nkernel = blackjax.nuts(log_prob, step_size=0.5)\nstate = kernel.init(jax.numpy.zeros(2))\n\nfor _ in range(1000):\n key, subkey = jax.random.split(key)\n state, info = kernel.step(subkey, state)\n```\n\n**Norn:**\n\n\n```ocaml\nlet metric = Norn.unit_metric 2 in\nlet kernel = Norn.nuts_kernel ~step_size:0.5 ~metric () in\nlet state = ref (kernel.init (Nx.zeros Nx.float64 [| 2 |]) log_prob) in\n\nfor _ = 1 to 1000 do\n let new_state, _info = kernel.step !state log_prob in\n state := new_state\ndone\n```\n\nBoth use a `{init; step}` pattern. The key difference: BlackJAX threads a\nPRNG key explicitly, while Norn uses Nx's RNG context (`Rng.run`).\n\n## Adaptation\n\n**BlackJAX:**\n\n```python\nwarmup = blackjax.window_adaptation(blackjax.nuts, log_prob)\nstate, kernel, _ = warmup.run(key, jax.numpy.zeros(2), 1000)\n```\n\n**Norn:**\n\n\n```ocaml\n(* Adaptation is built into sample/nuts/hmc *)\nlet result = Norn.nuts ~n:1000 ~num_warmup:500 log_prob init\n\n(* Or use sample for control over the kernel *)\nlet result =\n Norn.sample ~n:1000 ~num_warmup:500 log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~step_size ~metric ())\n```\n\nIn BlackJAX, adaptation is a separate step that returns a tuned kernel. In\nNorn, adaptation is integrated into `sample` -- it adapts step size and mass\nmatrix during warmup, then freezes them for sampling.\n\n## Samplers\n\n| PyMC / BlackJAX | Norn | Notes |\n|-----------------|------|-------|\n| `pm.sample()` (NUTS) | `Norn.nuts ~n log_prob init` | NUTS with adaptation |\n| `blackjax.nuts(log_prob, step_size)` | `Norn.nuts_kernel ~step_size ~metric ()` | NUTS kernel |\n| `blackjax.hmc(log_prob, step_size, ...)` | `Norn.hmc_kernel ~step_size ~metric ()` | HMC kernel |\n| `pm.sample(step=pm.HamiltonianMC(...))` | `Norn.hmc ~n log_prob init` | HMC with adaptation |\n\n## Integrators\n\n| BlackJAX | Norn | Notes |\n|----------|------|-------|\n| `blackjax.mcmc.integrators.velocity_verlet` | `Norn.leapfrog` | Default, 1 grad eval/step |\n| `blackjax.mcmc.integrators.mclachlan` | `Norn.mclachlan` | 2 grad evals/step |\n| `blackjax.mcmc.integrators.yoshida` | `Norn.yoshida` | 3 grad evals/step |\n\nUsage comparison:\n\n```python\n# BlackJAX\nkernel = blackjax.nuts(log_prob, step_size=0.5,\n integrator=blackjax.mcmc.integrators.mclachlan)\n```\n\n\n```ocaml\n(* Norn *)\nlet kernel =\n Norn.nuts_kernel ~integrator:Norn.mclachlan ~step_size:0.5 ~metric ()\n```\n\n## Metrics (Mass Matrix)\n\n| BlackJAX | Norn | Notes |\n|----------|------|-------|\n| `blackjax.mcmc.metrics.default_metric(jnp.ones(d))` | `Norn.unit_metric d` | Identity |\n| `blackjax.mcmc.metrics.default_metric(inv_mass_diag)` | `Norn.diagonal_metric inv_mass_diag` | Diagonal |\n| Dense metric via Cholesky | `Norn.dense_metric inv_mass_matrix` | Full matrix |\n\n## Diagnostics\n\n| PyMC / ArviZ | Norn | Notes |\n|--------------|------|-------|\n| `az.ess(trace)` | `Norn.ess samples` | Effective sample size |\n| `az.rhat(trace)` | `Norn.rhat chains` | Split R-hat |\n| `trace.sample_stats[\"diverging\"]` | `result.stats.num_divergent` | Divergence count |\n| `trace.sample_stats[\"accept\"]` | `result.stats.accept_rate` | Mean acceptance rate |\n| `trace.sample_stats[\"step_size\"]` | `result.stats.step_size` | Final step size |\n\n## State and Info\n\n**BlackJAX state:**\n\n```python\nstate.position # current sample\nstate.logdensity # log p(x)\nstate.logdensity_grad # grad log p(x)\n```\n\n**Norn state:**\n\n\n```ocaml\nstate.position (* Nx.float64_t, shape [dim] *)\nstate.log_density (* float *)\nstate.grad_log_density (* Nx.float64_t, shape [dim] *)\n```\n\n**BlackJAX info:**\n\n```python\ninfo.acceptance_rate\ninfo.is_divergent\ninfo.energy\ninfo.num_integration_steps\n```\n\n**Norn info:**\n\n\n```ocaml\ninfo.acceptance_rate (* float in [0, 1] *)\ninfo.is_divergent (* bool *)\ninfo.energy (* float *)\ninfo.num_integration_steps (* int *)\n```\n\n## Key Differences\n\n| Aspect | PyMC / BlackJAX | Norn |\n|--------|-----------------|------|\n| Language | Python / JAX | OCaml / Rune |\n| Model definition | Declarative (PyMC) or functional (BlackJAX) | Functional -- write `log_prob` directly |\n| Gradients | JAX autodiff | Rune autodiff |\n| PRNG | Explicit key splitting (JAX) | Scoped via `Nx.Rng.run` |\n| Adaptation | Separate step (BlackJAX) or automatic (PyMC) | Integrated into `sample` |\n| Mass matrix output | Diagonal or dense | `metric` record with `sample_momentum`, `kinetic_energy`, `scale` |\n| Multi-chain | Built-in (`chains` parameter) | Run multiple calls, combine with `rhat` |\n| Trace format | ArviZ InferenceData | `result` record with `samples` matrix |\n| Probabilistic DSL | Yes (PyMC) | No -- bring your own log-density |\n"
},
{
"path": "packages/norn/doc/dune",
"content": "(mdx\n (files *.md)\n (package norn)\n (libraries norn nx rune))\n"
},
{
"path": "packages/norn/doc/index.md",
"content": "# Norn\n\nNorn provides MCMC sampling with automatic gradients for OCaml. You supply an unnormalized log-density function and an initial position; Norn handles gradient computation via Rune, trajectory integration, and Stan-style window adaptation. One-line convenience functions cover common workflows, while the kernel API gives full control over integrators, metrics, and adaptation.\n\n## Features\n\n- **One-line sampling** -- `Norn.hmc` and `Norn.nuts` with automatic adaptation\n- **Configurable API** -- `Norn.sample` with custom kernels via `make_kernel`\n- **Automatic gradients** -- log-density gradients computed by Rune\n- **Symplectic integrators** -- `leapfrog`, `mclachlan`, `yoshida`\n- **Mass matrix metrics** -- `unit_metric`, `diagonal_metric`, `dense_metric`\n- **Stan-style adaptation** -- dual averaging for step size, Welford estimation for mass matrix\n- **Diagnostics** -- effective sample size (`ess`) and split R-hat (`rhat`)\n\n## Quick Start\n\n\n```ocaml\nopen Nx\n\nlet () =\n Rng.run ~seed:42 @@ fun () ->\n let f64 = Nx.float64 in\n\n (* Target: N([3; -1], I) *)\n let mu = Nx.create f64 [| 2 |] [| 3.0; -1.0 |] in\n let log_prob x =\n let d = Nx.sub x mu in\n Nx.mul_s (Nx.sum (Nx.square d)) (-0.5)\n in\n\n let init = Nx.zeros f64 [| 2 |] in\n let result = Norn.nuts ~n:1000 log_prob init in\n\n let mean = Nx.mean ~axes:[ 0 ] result.samples in\n Printf.printf \"posterior mean: %s\\n\" (Nx.data_to_string mean);\n Printf.printf \"accept rate: %.2f\\n\" result.stats.accept_rate;\n Printf.printf \"ESS: %s\\n\" (Nx.data_to_string (Norn.ess result.samples))\n```\n\n## Next Steps\n\n- [Getting Started](01-getting-started/) -- installation, first sampler, the kernel API\n- [Adaptation and Diagnostics](02-adaptation-and-diagnostics/) -- warmup windows, ESS, R-hat\n- [Advanced Usage](03-advanced-usage/) -- custom integrators, metrics, and monitoring\n- [PyMC Comparison](04-pymc-comparison/) -- mapping from Python's PyMC/BlackJAX to Norn\n"
},
{
"path": "packages/norn/examples/01-sampling-basics/README.md",
"content": "# `01-sampling-basics`\n\nYour first sampler. This example draws 1000 samples from a 2D correlated\nGaussian using NUTS and prints summary statistics to verify the chain recovered\nthe true distribution.\n\n```bash\ndune exec packages/norn/examples/01-sampling-basics/main.exe\n```\n\n## What You'll Learn\n\n- Defining a log-density function for MCMC\n- One-line sampling with `Norn.nuts`\n- Computing sample mean, variance, and covariance from the output\n- Reading basic diagnostics: ESS, acceptance rate, step size, divergences\n\n## Key Functions\n\n| Function | Purpose |\n| -------- | ------- |\n| `nuts` | Draw samples using the No-U-Turn Sampler with automatic adaptation |\n| `ess` | Effective sample size per parameter via autocorrelation |\n\n## How It Works\n\nThe target is a 2D Gaussian with mean `[2, -1]` and covariance `[[1, 0.8], [0.8, 2]]`.\nWe define `log_prob` as the unnormalized log-density (the Mahalanobis form), then\ncall `Norn.nuts ~n:1000 log_prob init`. NUTS handles warmup adaptation (step size\nand mass matrix) automatically.\n\n## Try It\n\n1. Increase `~n` to 5000 and observe ESS and variance estimates improve.\n2. Start from a bad initial point like `[100.0; 100.0]` -- warmup should still converge.\n3. Replace `Norn.nuts` with `Norn.hmc` and compare acceptance rates.\n\n## Next Steps\n\nContinue to [02-bayesian-regression](../02-bayesian-regression/) to see MCMC applied\nto a real inference problem.\n"
},
{
"path": "packages/norn/examples/01-sampling-basics/dune",
"content": "(executable\n (name main)\n (libraries nx rune norn))\n"
},
{
"path": "packages/norn/examples/01-sampling-basics/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Sample from a 2D correlated Gaussian using NUTS.\n\n Target distribution: N(mu, Sigma) with mu = [2.0; -1.0] Sigma = [[1.0, 0.8],\n [0.8, 2.0]]\n\n NUTS automatically adapts step size and trajectory length during warmup, so a\n single call to [Norn.nuts] is all you need. *)\n\nlet f64 = Nx.float64\n\n(* Target parameters *)\nlet mu = Nx.create f64 [| 2 |] [| 2.0; -1.0 |]\nlet sigma_inv = Nx.inv (Nx.create f64 [| 2; 2 |] [| 1.0; 0.8; 0.8; 2.0 |])\n\n(* Log-density of the target (unnormalized). *)\nlet log_prob x =\n let d = Nx.sub x mu in\n let dt = Nx.reshape [| 1; 2 |] d in\n let mahal = Nx.matmul (Nx.matmul dt sigma_inv) (Nx.reshape [| 2; 1 |] d) in\n Nx.mul_s (Nx.reshape [||] mahal) (-0.5)\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let init = Nx.zeros f64 [| 2 |] in\n let result = Norn.nuts ~n:1000 log_prob init in\n\n (* Sample mean *)\n let sample_mean = Nx.mean ~axes:[ 0 ] result.samples in\n Printf.printf \"--- 2D Correlated Gaussian (NUTS, 1000 samples) ---\\n\\n\";\n Printf.printf \"True mean: [%6.3f, %6.3f]\\n\" (Nx.item [ 0 ] mu)\n (Nx.item [ 1 ] mu);\n Printf.printf \"Sample mean: [%6.3f, %6.3f]\\n\"\n (Nx.item [ 0 ] sample_mean)\n (Nx.item [ 1 ] sample_mean);\n\n (* Sample variance *)\n let centered = Nx.sub result.samples sample_mean in\n let n = Float.of_int ((Nx.shape result.samples).(0) - 1) in\n let sample_cov =\n Nx.div_s (Nx.matmul (Nx.matrix_transpose centered) centered) n\n in\n Printf.printf \"\\nTrue var: [%6.3f, %6.3f]\\n\" 1.0 2.0;\n Printf.printf \"Sample var: [%6.3f, %6.3f]\\n\"\n (Nx.item [ 0; 0 ] sample_cov)\n (Nx.item [ 1; 1 ] sample_cov);\n Printf.printf \"True cov: %6.3f\\n\" 0.8;\n Printf.printf \"Sample cov: %6.3f\\n\" (Nx.item [ 0; 1 ] sample_cov);\n\n (* Diagnostics *)\n let e = Norn.ess result.samples in\n Printf.printf \"\\nESS: [%6.1f, %6.1f]\\n\" (Nx.item [ 0 ] e)\n (Nx.item [ 1 ] e);\n Printf.printf \"Accept rate: %.3f\\n\" result.stats.accept_rate;\n Printf.printf \"Step size: %.4f\\n\" result.stats.step_size;\n Printf.printf \"Divergent: %d\\n\" result.stats.num_divergent\n"
},
{
"path": "packages/norn/examples/02-bayesian-regression/README.md",
"content": "# `02-bayesian-regression`\n\nBayesian linear regression on synthetic data. Generates noisy observations from\n`y = 2x + 1`, defines a Gaussian likelihood with normal priors, and uses NUTS\nto infer the posterior over slope and intercept.\n\n```bash\ndune exec packages/norn/examples/02-bayesian-regression/main.exe\n```\n\n## What You'll Learn\n\n- Building a log-posterior from likelihood and prior\n- Interpreting posterior means and 95% credible intervals\n- Using the configurable `Norn.sample` API with `Norn.nuts_kernel`\n\n## Key Functions\n\n| Function | Purpose |\n| ------------- | ------------------------------------------------- |\n| `nuts` | One-line NUTS sampling with automatic adaptation |\n| `sample` | Configurable sampling with a user-provided kernel |\n| `nuts_kernel` | Construct a NUTS kernel with explicit parameters |\n| `ess` | Effective sample size per parameter |\n\n## How It Works\n\n1. Generate 50 data points from `y = 2x + 1 + N(0, 0.5)`.\n2. Define `log_posterior` as Gaussian log-likelihood plus `N(0, 10)` priors.\n3. Run `Norn.nuts ~n:2000` to draw posterior samples.\n4. Compute posterior means and 95% credible intervals from the samples.\n5. Re-run with `Norn.sample` + `Norn.nuts_kernel` to show the configurable API.\n\n## Try It\n\n1. Reduce `n_data` to 10 and observe wider credible intervals.\n2. Use a tighter prior `N(0, 1)` and see how it biases the posterior toward zero.\n3. Replace `Norn.nuts` with `Norn.hmc ~num_leapfrog:30` and compare.\n\n## Next Steps\n\nContinue to [03-diagnostics](../03-diagnostics/) to learn about multi-chain\nconvergence analysis with ESS and R-hat.\n"
},
{
"path": "packages/norn/examples/02-bayesian-regression/dune",
"content": "(executable\n (name main)\n (libraries nx rune norn))\n"
},
{
"path": "packages/norn/examples/02-bayesian-regression/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Bayesian linear regression: infer slope and intercept from noisy data.\n\n Model: y_i = slope * x_i + intercept + eps_i, eps_i ~ N(0, sigma^2)\n\n True parameters: slope = 2.0, intercept = 1.0, sigma = 0.5\n\n Priors: slope ~ N(0, 10) intercept ~ N(0, 10)\n\n We sample the posterior with NUTS and report credible intervals. *)\n\nlet f64 = Nx.float64\n\n(* Generate synthetic data: y = 2x + 1 + noise *)\nlet n_data = 50\nlet true_slope = 2.0\nlet true_intercept = 1.0\nlet noise_sigma = 0.5\n\nlet gen_data () =\n let x = Nx.linspace f64 (-2.0) 2.0 n_data in\n let noise = Nx.mul_s (Nx.randn f64 [| n_data |]) noise_sigma in\n let y =\n Nx.add (Nx.add (Nx.mul_s x true_slope) (Nx.scalar f64 true_intercept)) noise\n in\n (x, y)\n\n(* Log-posterior: Gaussian likelihood + normal prior. params = [slope;\n intercept] *)\nlet log_posterior x_data y_data params =\n let slope = Nx.slice [ I 0 ] params in\n let intercept = Nx.slice [ I 1 ] params in\n (* Predicted values *)\n let y_pred = Nx.add (Nx.mul x_data slope) intercept in\n let residuals = Nx.sub y_data y_pred in\n (* Log-likelihood: -0.5 * sum((y - y_pred)^2) / sigma^2 *)\n let ll =\n Nx.div_s\n (Nx.mul_s (Nx.sum (Nx.square residuals)) (-0.5))\n (noise_sigma *. noise_sigma)\n in\n (* Log-prior: N(0, 10) on each parameter *)\n let lp_slope = Nx.mul_s (Nx.square slope) (-0.5 /. 100.0) in\n let lp_intercept = Nx.mul_s (Nx.square intercept) (-0.5 /. 100.0) in\n Nx.add ll (Nx.add lp_slope lp_intercept)\n\nlet percentile samples frac =\n let n = (Nx.shape samples).(0) in\n let sorted, _ = Nx.sort samples in\n let idx = Float.to_int (frac *. Float.of_int (n - 1)) in\n Nx.item [ idx ] sorted\n\nlet () =\n Nx.Rng.run ~seed:42 @@ fun () ->\n let x_data, y_data = gen_data () in\n let init = Nx.zeros f64 [| 2 |] in\n let log_prob = log_posterior x_data y_data in\n let result = Norn.nuts ~n:2000 ~num_warmup:1000 log_prob init in\n\n Printf.printf \"--- Bayesian Linear Regression (NUTS, 2000 samples) ---\\n\\n\";\n Printf.printf \"True: slope = %.2f, intercept = %.2f\\n\" true_slope\n true_intercept;\n\n let sample_mean = Nx.mean ~axes:[ 0 ] result.samples in\n Printf.printf \"Posterior: slope = %.3f, intercept = %.3f\\n\"\n (Nx.item [ 0 ] sample_mean)\n (Nx.item [ 1 ] sample_mean);\n\n (* 95%% credible intervals *)\n Printf.printf \"\\n95%% credible intervals:\\n\";\n let slope_samples = Nx.slice [ A; I 0 ] result.samples in\n let intercept_samples = Nx.slice [ A; I 1 ] result.samples in\n Printf.printf \" slope: [%.3f, %.3f]\\n\"\n (percentile slope_samples 0.025)\n (percentile slope_samples 0.975);\n Printf.printf \" intercept: [%.3f, %.3f]\\n\"\n (percentile intercept_samples 0.025)\n (percentile intercept_samples 0.975);\n\n (* Diagnostics *)\n let e = Norn.ess result.samples in\n Printf.printf \"\\nESS: [%.1f, %.1f]\\n\" (Nx.item [ 0 ] e)\n (Nx.item [ 1 ] e);\n Printf.printf \"Accept rate: %.3f\\n\" result.stats.accept_rate;\n Printf.printf \"Step size: %.4f\\n\" result.stats.step_size;\n Printf.printf \"Divergent: %d\\n\" result.stats.num_divergent;\n\n (* Also demonstrate the configurable API *)\n Printf.printf \"\\n--- Same model with configurable sample API ---\\n\";\n let result2 =\n Norn.sample ~n:1000 ~num_warmup:500 log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~step_size ~metric ())\n in\n let mean2 = Nx.mean ~axes:[ 0 ] result2.samples in\n Printf.printf \"Posterior: slope = %.3f, intercept = %.3f\\n\"\n (Nx.item [ 0 ] mean2) (Nx.item [ 1 ] mean2);\n Printf.printf \"Accept rate: %.3f\\n\" result2.stats.accept_rate\n"
},
{
"path": "packages/norn/examples/03-diagnostics/README.md",
"content": "# `03-diagnostics`\n\nMulti-chain convergence diagnostics. Runs 4 independent NUTS chains on a 3D\nGaussian target, then computes ESS and split R-hat to verify that the chains\nhave converged and mixed.\n\n```bash\ndune exec packages/norn/examples/03-diagnostics/main.exe\n```\n\n## What You'll Learn\n\n- Running multiple chains with different seeds\n- Computing effective sample size (ESS) per chain\n- Computing split R-hat across chains for convergence assessment\n- Interpreting diagnostic thresholds (ESS > 100, R-hat < 1.01)\n\n## Key Functions\n\n| Function | Purpose |\n| -------- | ------- |\n| `nuts` | Draw samples with automatic adaptation |\n| `ess` | Effective sample size via initial monotone sequence estimator |\n| `rhat` | Split R-hat convergence diagnostic across chains |\n\n## How It Works\n\n1. Define a 3D Gaussian target with different scales per dimension (var = 1, 4, 0.25).\n2. Run 4 independent NUTS chains, each with a different random seed.\n3. Report per-chain acceptance rate, step size, and divergence count.\n4. Compute ESS for each chain and R-hat across chains.\n5. Pool all chains for a final posterior summary.\n\n## Interpreting the Output\n\n- **ESS**: The number of effectively independent samples. If ESS is much lower\n than the actual sample count, the chain has high autocorrelation.\n- **R-hat**: Measures between-chain vs within-chain variance. Values close to\n 1.0 mean the chains agree. Above 1.01 suggests incomplete mixing.\n\n## Try It\n\n1. Reduce `n_samples` to 50 and observe R-hat increase above 1.01.\n2. Use a highly correlated target and see ESS drop.\n3. Try `Norn.hmc` instead of `Norn.nuts` and compare ESS.\n\n## Further Reading\n\n- [Sampling Basics](../01-sampling-basics/) -- single-chain sampling\n- [Bayesian Regression](../02-bayesian-regression/) -- a real inference problem\n"
},
{
"path": "packages/norn/examples/03-diagnostics/dune",
"content": "(executable\n (name main)\n (libraries nx rune norn))\n"
},
{
"path": "packages/norn/examples/03-diagnostics/main.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Multi-chain convergence diagnostics.\n\n Run 4 chains on a 3D target distribution, then compute ESS and R-hat to\n assess whether the chains have converged and mixed well.\n\n Guidelines: - ESS > 100 per parameter for reliable estimates - R-hat < 1.01\n indicates convergence across chains *)\n\nlet f64 = Nx.float64\n\n(* Target: 3D Gaussian with different scales per dimension. mu = [1, -2, 0.5]\n sigma = diag([1, 4, 0.25]) *)\nlet mu = Nx.create f64 [| 3 |] [| 1.0; -2.0; 0.5 |]\nlet inv_var = Nx.create f64 [| 3 |] [| 1.0; 0.25; 4.0 |]\n(* 1/sigma^2 for each dim *)\n\nlet log_prob x =\n let d = Nx.sub x mu in\n Nx.mul_s (Nx.sum (Nx.mul (Nx.square d) inv_var)) (-0.5)\n\nlet dim = 3\nlet n_chains = 4\nlet n_samples = 1000\nlet param_names = [| \"x0\"; \"x1\"; \"x2\" |]\n\nlet () =\n Printf.printf \"--- Multi-Chain Diagnostics (%d chains x %d samples) ---\\n\\n\"\n n_chains n_samples;\n\n (* Run chains with different seeds *)\n let chains =\n Array.init n_chains (fun i ->\n Nx.Rng.run ~seed:(i + 1) @@ fun () ->\n let init = Nx.zeros f64 [| dim |] in\n Norn.nuts ~n:n_samples ~num_warmup:500 log_prob init)\n in\n\n (* Per-chain summary *)\n Printf.printf \"Per-chain summary:\\n\";\n Printf.printf \" %-8s %-12s %-12s %-8s\\n\" \"Chain\" \"Accept Rate\" \"Step Size\"\n \"Diverg.\";\n Array.iteri\n (fun i r ->\n Printf.printf \" %-8d %-12.3f %-12.4f %-8d\\n\" (i + 1)\n r.Norn.stats.accept_rate r.stats.step_size r.stats.num_divergent)\n chains;\n\n (* Per-chain ESS *)\n Printf.printf \"\\nEffective Sample Size (ESS) per chain:\\n\";\n Printf.printf \" %-8s\" \"Chain\";\n Array.iter (fun name -> Printf.printf \" %-8s\" name) param_names;\n Printf.printf \"\\n\";\n Array.iteri\n (fun i r ->\n let e = Norn.ess r.Norn.samples in\n Printf.printf \" %-8d\" (i + 1);\n for d = 0 to dim - 1 do\n Printf.printf \" %-8.1f\" (Nx.item [ d ] e)\n done;\n Printf.printf \"\\n\")\n chains;\n\n (* R-hat across chains *)\n let chain_samples = Array.map (fun r -> r.Norn.samples) chains in\n let r = Norn.rhat chain_samples in\n Printf.printf \"\\nSplit R-hat (target: < 1.01):\\n\";\n for d = 0 to dim - 1 do\n let rv = Nx.item [ d ] r in\n let status = if rv < 1.01 then \"OK\" else \"WARNING\" in\n Printf.printf \" %s: %.4f [%s]\\n\" param_names.(d) rv status\n done;\n\n (* Grand summary *)\n let all_converged = ref true in\n for d = 0 to dim - 1 do\n if Nx.item [ d ] r >= 1.01 then all_converged := false\n done;\n Printf.printf \"\\nConvergence: %s\\n\"\n (if !all_converged then \"All parameters converged (R-hat < 1.01)\"\n else\n \"Some parameters have not converged -- increase samples or check model\");\n\n (* Pooled posterior summary *)\n Printf.printf \"\\nPooled posterior (all chains):\\n\";\n Printf.printf \" %-8s %-10s %-10s %-10s\\n\" \"Param\" \"True\" \"Mean\" \"Std\";\n let all_samples =\n Nx.concatenate ~axis:0\n (Array.to_list (Array.map (fun r -> r.Norn.samples) chains))\n in\n let pooled_mean = Nx.mean ~axes:[ 0 ] all_samples in\n let pooled_centered = Nx.sub all_samples pooled_mean in\n let nf = Float.of_int ((Nx.shape all_samples).(0) - 1) in\n let pooled_var =\n Nx.div_s (Nx.sum ~axes:[ 0 ] (Nx.square pooled_centered)) nf\n in\n let pooled_std = Nx.sqrt pooled_var in\n for d = 0 to dim - 1 do\n Printf.printf \" %-8s %-10.3f %-10.3f %-10.3f\\n\" param_names.(d)\n (Nx.item [ d ] mu)\n (Nx.item [ d ] pooled_mean)\n (Nx.item [ d ] pooled_std)\n done\n"
},
{
"path": "packages/norn/examples/README.md",
"content": "# Norn Examples\n\nLearn Norn through progressively complex examples. Start with `01-sampling-basics`\nand work through the numbered examples in order.\n\n## Examples\n\n| Example | Concept | Key Functions |\n|---------|---------|---------------|\n| [`01-sampling-basics`](./01-sampling-basics/) | Sample from a correlated Gaussian with NUTS | `nuts`, `ess` |\n| [`02-bayesian-regression`](./02-bayesian-regression/) | Bayesian linear regression with posterior inference | `nuts`, `sample`, `nuts_kernel` |\n| [`03-diagnostics`](./03-diagnostics/) | Multi-chain convergence diagnostics | `nuts`, `ess`, `rhat` |\n\n## Running Examples\n\nAll examples can be run with:\n\n```bash\ndune exec packages/norn/examples//main.exe\n```\n\nFor example:\n\n```bash\ndune exec packages/norn/examples/01-sampling-basics/main.exe\n```\n\n## Quick Reference\n\n### One-Line Sampling\n\n```ocaml\nlet result =\n Nx.Rng.run ~seed:42 @@ fun () ->\n Norn.nuts ~n:1000 log_prob (Nx.zeros Nx.float64 [| dim |])\n```\n\n### Configurable Sampling\n\n```ocaml\nlet result =\n Norn.sample ~n:1000 ~num_warmup:500 log_prob init\n (fun ~step_size ~metric ->\n Norn.nuts_kernel ~step_size ~metric ())\n```\n\n### Convergence Diagnostics\n\n```ocaml\nlet ess = Norn.ess result.samples in\nlet rhat = Norn.rhat [| chain1.samples; chain2.samples |]\n```\n"
},
{
"path": "packages/norn/lib/adapt.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Step-size adaptation via dual averaging (Nesterov 2009). *)\n\nlet f64 = Nx.float64\n\ntype step_size = {\n target_accept : float;\n mu : float;\n log_eps : float;\n log_eps_bar : float;\n h_bar : float;\n count : int;\n}\n\nlet step_size_init ?(target_accept = 0.65) eps =\n {\n target_accept;\n mu = Float.log (10.0 *. eps);\n log_eps = Float.log eps;\n log_eps_bar = 0.0;\n h_bar = 0.0;\n count = 0;\n }\n\nlet step_size_update ss ~acceptance_rate =\n let gamma = 0.05 in\n let t0 = 10.0 in\n let kappa = 0.75 in\n let m = Float.of_int (ss.count + 1) in\n let w = 1.0 /. (m +. t0) in\n let h_bar =\n ((1.0 -. w) *. ss.h_bar) +. (w *. (ss.target_accept -. acceptance_rate))\n in\n let log_eps = ss.mu -. (Float.sqrt m /. gamma *. h_bar) in\n let m_pow = m ** -.kappa in\n (* BlackJAX: log_x_avg uses PREVIOUS log_x, not the newly computed one. *)\n let log_eps_bar =\n (m_pow *. ss.log_eps) +. ((1.0 -. m_pow) *. ss.log_eps_bar)\n in\n { ss with h_bar; log_eps; log_eps_bar; count = ss.count + 1 }\n\nlet step_size_current ss = Float.exp ss.log_eps\nlet step_size_final ss = Float.exp ss.log_eps_bar\n\n(* Mass-matrix adaptation via Welford's online algorithm. *)\n\ntype mass_matrix = {\n dim : int;\n count : int;\n mean : Nx.float64_t;\n m2 : Nx.float64_t;\n}\n\nlet mass_matrix_init dim =\n { dim; count = 0; mean = Nx.zeros f64 [| dim |]; m2 = Nx.zeros f64 [| dim |] }\n\nlet mass_matrix_update mm position =\n let count = mm.count + 1 in\n let delta = Nx.sub position mm.mean in\n let mean = Nx.add mm.mean (Nx.div_s delta (Float.of_int count)) in\n let delta2 = Nx.sub position mean in\n let m2 = Nx.add mm.m2 (Nx.mul delta delta2) in\n { mm with count; mean; m2 }\n\nlet mass_matrix_inv_diag mm =\n if mm.count < 2 then None\n else\n let n = Float.of_int mm.count in\n let variance = Nx.div_s mm.m2 (n -. 1.0) in\n let w = n /. (n +. 5.0) in\n let shrinkage = 1e-3 *. 5.0 /. (n +. 5.0) in\n Some (Nx.add_s (Nx.mul_s variance w) shrinkage)\n\nlet mass_matrix_reset mm =\n {\n mm with\n count = 0;\n mean = Nx.zeros f64 [| mm.dim |];\n m2 = Nx.zeros f64 [| mm.dim |];\n }\n\nlet step_size_reset ss =\n step_size_init ~target_accept:ss.target_accept (step_size_final ss)\n\n(* Window adaptation schedule (Stan warmup).\n\n Three phases: - Fast (initial buffer): adapt step size only. - Slow (doubling\n windows): adapt step size + mass matrix. At each window boundary the mass\n matrix is finalized (regularized) and both estimators are reset. - Fast\n (final buffer): adapt step size only with the final mass matrix. *)\n\ntype warmup_action = Fast | Slow | Slow_end\n\nlet build_schedule num_warmup =\n if num_warmup < 20 then Array.make num_warmup Fast\n else\n let initial_buffer, final_buffer, first_window =\n if 75 + 50 + 25 > num_warmup then\n let ib = Float.to_int (0.15 *. Float.of_int num_warmup) in\n let fb = Float.to_int (0.10 *. Float.of_int num_warmup) in\n (ib, fb, num_warmup - ib - fb)\n else (75, 50, 25)\n in\n let schedule = Array.make num_warmup Fast in\n let slow_end_pos = num_warmup - final_buffer in\n let pos = ref initial_buffer in\n let window_size = ref first_window in\n while !pos < slow_end_pos do\n let end_pos =\n if !pos + (3 * !window_size) <= slow_end_pos then !pos + !window_size\n else slow_end_pos\n in\n for j = !pos to end_pos - 1 do\n schedule.(j) <- (if j = end_pos - 1 then Slow_end else Slow)\n done;\n pos := end_pos;\n window_size := !window_size * 2\n done;\n schedule\n"
},
{
"path": "packages/norn/lib/dune",
"content": "(library\n (name norn)\n (public_name norn)\n (private_modules adapt internal nuts)\n (libraries nx rune))\n"
},
{
"path": "packages/norn/lib/internal.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ntype state = {\n position : Nx.float64_t;\n log_density : float;\n grad_log_density : Nx.float64_t;\n}\n\ntype info = {\n acceptance_rate : float;\n is_divergent : bool;\n energy : float;\n num_integration_steps : int;\n}\n\ntype kernel = {\n init : Nx.float64_t -> (Nx.float64_t -> Nx.float64_t) -> state;\n step : state -> (Nx.float64_t -> Nx.float64_t) -> state * info;\n}\n\ntype integrator =\n (Nx.float64_t -> Nx.float64_t) ->\n Nx.float64_t ->\n Nx.float64_t ->\n Nx.float64_t ->\n (Nx.float64_t -> float * Nx.float64_t) ->\n float ->\n Nx.float64_t * Nx.float64_t * float * Nx.float64_t\n\ntype metric = {\n sample_momentum : int -> Nx.float64_t;\n kinetic_energy : Nx.float64_t -> float;\n scale : Nx.float64_t -> Nx.float64_t;\n is_turning : Nx.float64_t -> Nx.float64_t -> Nx.float64_t -> bool;\n}\n\ntype stats = { accept_rate : float; step_size : float; num_divergent : int }\n\ntype result = {\n samples : Nx.float64_t;\n log_densities : Nx.float64_t;\n stats : stats;\n}\n"
},
{
"path": "packages/norn/lib/norn.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ninclude Internal\n\nlet f64 = Nx.float64\n\n(* Integrators *)\n\nlet leapfrog kinetic_energy_grad q p g grad_log_prob eps =\n let p = Nx.add p (Nx.mul_s g (eps /. 2.0)) in\n let q = Nx.add q (Nx.mul_s (kinetic_energy_grad p) eps) in\n let lp, g = grad_log_prob q in\n let p = Nx.add p (Nx.mul_s g (eps /. 2.0)) in\n (q, p, lp, g)\n\nlet palindromic ~momentum_coeffs ~position_coeffs kinetic_energy_grad q p g\n grad_log_prob eps =\n let n_pos = Array.length position_coeffs in\n let q = ref q in\n let p = ref (Nx.add p (Nx.mul_s g (momentum_coeffs.(0) *. eps))) in\n let lp = ref 0.0 in\n let g = ref g in\n for i = 0 to n_pos - 1 do\n q :=\n Nx.add !q (Nx.mul_s (kinetic_energy_grad !p) (position_coeffs.(i) *. eps));\n let lp', g' = grad_log_prob !q in\n lp := lp';\n g := g';\n p := Nx.add !p (Nx.mul_s !g (momentum_coeffs.(i + 1) *. eps))\n done;\n (!q, !p, !lp, !g)\n\nlet mclachlan =\n let l = 0.1932093174209856 in\n palindromic\n ~momentum_coeffs:[| l; 1.0 -. (2.0 *. l); l |]\n ~position_coeffs:[| 0.5; 0.5 |]\n\nlet yoshida =\n let cbrt2 = 2.0 ** (1.0 /. 3.0) in\n let w1 = 1.0 /. (2.0 -. cbrt2) in\n let w0 = -.cbrt2 /. (2.0 -. cbrt2) in\n palindromic\n ~momentum_coeffs:\n [| w1 /. 2.0; (w1 +. w0) /. 2.0; (w0 +. w1) /. 2.0; w1 /. 2.0 |]\n ~position_coeffs:[| w1; w0; w1 |]\n\n(* Metrics *)\n\nlet euclidean_is_turning scale left_p right_p momentum_sum =\n let rho = Nx.sub momentum_sum (Nx.div_s (Nx.add left_p right_p) 2.0) in\n Nx.item [] (Nx.vdot (scale left_p) rho) <= 0.0\n || Nx.item [] (Nx.vdot (scale right_p) rho) <= 0.0\n\nlet unit_metric dim =\n {\n sample_momentum = (fun _dim -> Nx.randn f64 [| dim |]);\n kinetic_energy = (fun p -> 0.5 *. Nx.item [] (Nx.sum (Nx.square p)));\n scale = Fun.id;\n is_turning = euclidean_is_turning Fun.id;\n }\n\nlet diagonal_metric inv_mass_diag =\n let mass_diag = Nx.recip inv_mass_diag in\n let sqrt_mass = Nx.sqrt mass_diag in\n let scale v = Nx.mul v inv_mass_diag in\n {\n sample_momentum = (fun dim -> Nx.mul (Nx.randn f64 [| dim |]) sqrt_mass);\n kinetic_energy =\n (fun p -> 0.5 *. Nx.item [] (Nx.sum (Nx.mul (Nx.square p) inv_mass_diag)));\n scale;\n is_turning = euclidean_is_turning scale;\n }\n\nlet dense_metric inv_mass_matrix =\n let dim = (Nx.shape inv_mass_matrix).(0) in\n let mass_matrix = Nx.inv inv_mass_matrix in\n let chol = Nx.cholesky mass_matrix in\n let scale v =\n let v_col = Nx.reshape [| dim; 1 |] v in\n Nx.reshape [| dim |] (Nx.matmul inv_mass_matrix v_col)\n in\n {\n sample_momentum =\n (fun _dim ->\n let z = Nx.randn f64 [| dim; 1 |] in\n Nx.reshape [| dim |] (Nx.matmul chol z));\n kinetic_energy =\n (fun p ->\n let p_col = Nx.reshape [| dim; 1 |] p in\n 0.5\n *. Nx.item []\n (Nx.matmul\n (Nx.matrix_transpose p_col)\n (Nx.matmul inv_mass_matrix p_col)));\n scale;\n is_turning = euclidean_is_turning scale;\n }\n\n(* Kernels *)\n\nlet grad_log_prob log_density_fn q =\n let lp, g = Rune.value_and_grad log_density_fn q in\n (Nx.item [] lp, g)\n\nlet init_state position log_density_fn =\n let lp, g = Rune.value_and_grad log_density_fn position in\n { position; log_density = Nx.item [] lp; grad_log_density = g }\n\nlet hmc_kernel ?(integrator : integrator = leapfrog) ?(num_leapfrog = 20)\n ~step_size ~(metric : metric) () =\n let step (state : state) log_density_fn =\n let dim = Nx.numel state.position in\n let glp = grad_log_prob log_density_fn in\n let p0 = metric.sample_momentum dim in\n let ke_current = metric.kinetic_energy p0 in\n let q = ref state.position in\n let p = ref p0 in\n let lp = ref state.log_density in\n let g = ref state.grad_log_density in\n for _ = 1 to num_leapfrog do\n let q', p', lp', g' = integrator metric.scale !q !p !g glp step_size in\n q := q';\n p := p';\n lp := lp';\n g := g'\n done;\n let ke_proposed = metric.kinetic_energy !p in\n let delta = !lp -. state.log_density -. (ke_proposed -. ke_current) in\n let log_accept = if Float.is_nan delta then Float.neg_infinity else delta in\n let acceptance_rate = Float.min 1.0 (Float.exp log_accept) in\n let accepted = Float.log (Nx.item [] (Nx.rand f64 [||])) < log_accept in\n let new_state =\n if accepted then\n { position = !q; log_density = !lp; grad_log_density = !g }\n else state\n in\n let info =\n {\n acceptance_rate;\n is_divergent = Float.abs (ke_proposed -. ke_current) > 1000.0;\n energy = -. !lp +. ke_proposed;\n num_integration_steps = num_leapfrog;\n }\n in\n (new_state, info)\n in\n { init = init_state; step }\n\nlet nuts_kernel ?(integrator : integrator = leapfrog) ?(max_depth = 10)\n ~step_size ~(metric : metric) () =\n let step state log_density_fn =\n Nuts.step integrator metric step_size max_depth state log_density_fn\n in\n { init = init_state; step }\n\n(* Sampling *)\n\nlet metric_of_mass_matrix dim mm =\n match Adapt.mass_matrix_inv_diag mm with\n | None -> unit_metric dim\n | Some inv_mass_diag -> diagonal_metric inv_mass_diag\n\nlet sample ?(step_size = 0.01) ?(target_accept = 0.65) ?num_warmup ?report ~n\n log_density_fn init make_kernel =\n let num_warmup = match num_warmup with Some w -> w | None -> n / 2 in\n let dim = Nx.numel init in\n let schedule = Adapt.build_schedule num_warmup in\n let met = ref (unit_metric dim) in\n let kern = ref (make_kernel ~step_size ~metric:!met) in\n let state = ref (!kern.init init log_density_fn) in\n let ss = ref (Adapt.step_size_init ~target_accept step_size) in\n let mm = ref (Adapt.mass_matrix_init dim) in\n for i = 1 to num_warmup do\n let eps = Adapt.step_size_current !ss in\n kern := make_kernel ~step_size:eps ~metric:!met;\n let new_state, info = !kern.step !state log_density_fn in\n state := new_state;\n (match schedule.(i - 1) with\n | Adapt.Fast ->\n ss := Adapt.step_size_update !ss ~acceptance_rate:info.acceptance_rate\n | Adapt.Slow ->\n ss := Adapt.step_size_update !ss ~acceptance_rate:info.acceptance_rate;\n mm := Adapt.mass_matrix_update !mm new_state.position\n | Adapt.Slow_end ->\n ss := Adapt.step_size_update !ss ~acceptance_rate:info.acceptance_rate;\n mm := Adapt.mass_matrix_update !mm new_state.position;\n met := metric_of_mass_matrix dim !mm;\n mm := Adapt.mass_matrix_reset !mm;\n ss := Adapt.step_size_reset !ss);\n match report with\n | Some f -> f ~step:(-(num_warmup - i + 1)) new_state info\n | None -> ()\n done;\n let final_step_size = Adapt.step_size_final !ss in\n kern := make_kernel ~step_size:final_step_size ~metric:!met;\n let samples = Nx.zeros f64 [| n; dim |] in\n let log_densities = Nx.zeros f64 [| n |] in\n let total_accept = ref 0.0 in\n let num_divergent = ref 0 in\n for i = 0 to n - 1 do\n let new_state, info = !kern.step !state log_density_fn in\n state := new_state;\n total_accept := !total_accept +. info.acceptance_rate;\n if info.is_divergent then incr num_divergent;\n Nx.set_slice [ I i ] samples new_state.position;\n Nx.set_item [ i ] new_state.log_density log_densities;\n match report with Some f -> f ~step:i new_state info | None -> ()\n done;\n {\n samples;\n log_densities;\n stats =\n {\n accept_rate = !total_accept /. Float.of_int n;\n step_size = final_step_size;\n num_divergent = !num_divergent;\n };\n }\n\nlet hmc ?(step_size = 0.01) ?(target_accept = 0.65) ?num_leapfrog ?num_warmup ~n\n log_prob init =\n sample ~step_size ~target_accept ?num_warmup ~n log_prob init\n (fun ~step_size ~metric ->\n hmc_kernel ?integrator:None ?num_leapfrog ~step_size ~metric ())\n\nlet nuts ?(step_size = 0.01) ?(target_accept = 0.80) ?max_depth ?num_warmup ~n\n log_prob init =\n sample ~step_size ~target_accept ?num_warmup ~n log_prob init\n (fun ~step_size ~metric ->\n nuts_kernel ?integrator:None ?max_depth ~step_size ~metric ())\n\n(* Diagnostics *)\n\nlet autocorr samples =\n let n = (Nx.shape samples).(0) in\n let dim = (Nx.shape samples).(1) in\n let mean = Nx.mean ~axes:[ 0 ] samples in\n let centered = Nx.sub samples mean in\n let max_lag = n / 2 in\n let acf = Nx.zeros f64 [| max_lag; dim |] in\n for d = 0 to dim - 1 do\n let col = Nx.slice [ A; I d ] centered in\n let v = ref 0.0 in\n for i = 0 to n - 1 do\n let x = Nx.item [ i ] col in\n v := !v +. (x *. x)\n done;\n v := !v /. Float.of_int n;\n for lag = 0 to max_lag - 1 do\n let c = ref 0.0 in\n for i = 0 to n - 1 - lag do\n c := !c +. (Nx.item [ i ] col *. Nx.item [ i + lag ] col)\n done;\n Nx.set_item [ lag; d ] (!c /. (Float.of_int n *. !v)) acf\n done\n done;\n acf\n\nlet ess samples =\n let n = (Nx.shape samples).(0) in\n let dim = (Nx.shape samples).(1) in\n let acf = autocorr samples in\n let max_lag = n / 2 in\n let result = Nx.zeros f64 [| dim |] in\n for d = 0 to dim - 1 do\n let tau = ref 1.0 in\n let lag = ref 1 in\n let stop = ref false in\n while !lag < max_lag - 1 && not !stop do\n let rho1 = Nx.item [ !lag; d ] acf in\n let rho2 = Nx.item [ !lag + 1; d ] acf in\n if rho1 +. rho2 < 0.0 then stop := true\n else begin\n tau := !tau +. (2.0 *. rho1);\n incr lag\n end\n done;\n Nx.set_item [ d ] (Float.of_int n /. !tau) result\n done;\n result\n\nlet rhat chains =\n let m = Array.length chains in\n let n = (Nx.shape chains.(0)).(0) in\n let dim = (Nx.shape chains.(0)).(1) in\n let half = n / 2 in\n let split_chains = Array.make (2 * m) chains.(0) in\n for i = 0 to m - 1 do\n split_chains.(2 * i) <- Nx.slice [ R (0, half - 1) ] chains.(i);\n split_chains.((2 * i) + 1) <- Nx.slice [ R (half, n - 1) ] chains.(i)\n done;\n let nf = Float.of_int half in\n let mf = Float.of_int (2 * m) in\n let chain_means = Array.map (Nx.mean ~axes:[ 0 ]) split_chains in\n let grand_mean =\n Array.fold_left Nx.add (Nx.zeros f64 [| dim |]) chain_means |> fun s ->\n Nx.div_s s mf\n in\n let b =\n Array.fold_left\n (fun acc cm ->\n let diff = Nx.sub cm grand_mean in\n Nx.add acc (Nx.square diff))\n (Nx.zeros f64 [| dim |]) chain_means\n |> fun s -> Nx.mul_s s (nf /. (mf -. 1.0))\n in\n let w =\n Array.fold_left\n (fun acc chain ->\n let cm = Nx.mean ~axes:[ 0 ] chain in\n let centered = Nx.sub chain cm in\n let s2 =\n Nx.div_s (Nx.sum ~axes:[ 0 ] (Nx.square centered)) (nf -. 1.0)\n in\n Nx.add acc s2)\n (Nx.zeros f64 [| dim |]) split_chains\n |> fun s -> Nx.div_s s mf\n in\n let var_hat = Nx.add (Nx.mul_s w ((nf -. 1.0) /. nf)) (Nx.div_s b nf) in\n Nx.sqrt (Nx.div var_hat w)\n"
},
{
"path": "packages/norn/lib/norn.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(** MCMC sampling with automatic gradients.\n\n Norn provides Markov chain Monte Carlo samplers that leverage {!Rune}'s\n automatic differentiation. The core abstraction is the {!type-kernel}: a\n composable [{init; step}] record that any algorithm produces and any\n sampling loop consumes.\n\n {b Quick start.}\n {[\n let result = Norn.nuts ~n:1000 log_prob (Nx.zeros Nx.float64 [| dim |])\n ]}\n\n For configured usage, construct a kernel and pass it to {!sample}:\n {[\n let result =\n Norn.sample ~n:1000 log_prob init (fun ~step_size ~metric ->\n Norn.nuts_kernel ~step_size ~metric ())\n ]} *)\n\n(** {1:types Types} *)\n\ntype state = {\n position : Nx.float64_t; (** Current sample, shape [[dim]]. *)\n log_density : float; (** Log-density at {!position}. *)\n grad_log_density : Nx.float64_t;\n (** Gradient of log-density at {!position}, shape [[dim]]. *)\n}\n(** The type for sampler states. Shared across all gradient-based kernels. *)\n\ntype info = {\n acceptance_rate : float;\n (** Metropolis acceptance probability in \\[0, 1\\]. *)\n is_divergent : bool; (** [true] when the energy error exceeds 1000. *)\n energy : float; (** Total Hamiltonian energy of the proposal. *)\n num_integration_steps : int; (** Leapfrog steps taken this transition. *)\n}\n(** The type for per-step diagnostics. *)\n\ntype kernel = {\n init : Nx.float64_t -> (Nx.float64_t -> Nx.float64_t) -> state;\n (** [init position log_density_fn] is the initial state at [position]. *)\n step : state -> (Nx.float64_t -> Nx.float64_t) -> state * info;\n (** [step state log_density_fn] is [(new_state, info)]. *)\n}\n(** The type for sampling kernels. Constructed by {!hmc_kernel}, {!nuts_kernel},\n etc. The [log_density_fn] argument is not baked in so the same kernel can be\n reused with different targets (e.g. tempering). *)\n\n(** {1:integrators Integrators} *)\n\ntype integrator =\n (Nx.float64_t -> Nx.float64_t) ->\n Nx.float64_t ->\n Nx.float64_t ->\n Nx.float64_t ->\n (Nx.float64_t -> float * Nx.float64_t) ->\n float ->\n Nx.float64_t * Nx.float64_t * float * Nx.float64_t\n(** The type for symplectic integrators.\n [integrator kinetic_energy_grad position momentum gradient grad_log_prob\n step_size] is [(new_pos, new_mom, new_log_density, new_grad)].\n\n [kinetic_energy_grad] is [M{^-1} p], the gradient of the kinetic energy with\n respect to momentum. For unit metric this is the identity. The kernel\n provides it from {!type-metric}[.scale]. *)\n\nval leapfrog : integrator\n(** [leapfrog] is the velocity Verlet integrator (second-order symplectic). *)\n\nval mclachlan : integrator\n(** [mclachlan] is McLachlan's two-stage integrator. Higher acceptance rates\n than {!leapfrog} on challenging posteriors (McLachlan 1995). Two gradient\n evaluations per step. *)\n\nval yoshida : integrator\n(** [yoshida] is Yoshida's fourth-order symplectic integrator. More accurate\n than {!leapfrog} at the cost of three gradient evaluations per step. *)\n\n(** {1:metrics Metrics} *)\n\ntype metric = {\n sample_momentum : int -> Nx.float64_t;\n (** [sample_momentum dim] draws momentum from the kinetic energy\n distribution. *)\n kinetic_energy : Nx.float64_t -> float;\n (** [kinetic_energy p] is [0.5 * p{^T} M{^-1} p]. *)\n scale : Nx.float64_t -> Nx.float64_t; (** [scale v] is [M{^-1} v]. *)\n is_turning : Nx.float64_t -> Nx.float64_t -> Nx.float64_t -> bool;\n (** [is_turning left_p right_p momentum_sum] is the U-turn criterion for\n NUTS trajectory termination. *)\n}\n(** The type for mass matrix metrics. Defines the geometry of the sampling\n space. *)\n\nval unit_metric : int -> metric\n(** [unit_metric dim] is the identity metric. Momentum sampled from [N(0, I)].\n*)\n\nval diagonal_metric : Nx.float64_t -> metric\n(** [diagonal_metric inv_mass_diag] is a diagonal metric with the given inverse\n mass diagonal. *)\n\nval dense_metric : Nx.float64_t -> metric\n(** [dense_metric inv_mass_matrix] is a dense metric with the given inverse mass\n matrix. Uses Cholesky decomposition for momentum sampling. *)\n\n(** {1:kernels Kernels} *)\n\nval hmc_kernel :\n ?integrator:integrator ->\n ?num_leapfrog:int ->\n step_size:float ->\n metric:metric ->\n unit ->\n kernel\n(** [hmc_kernel ~step_size ~metric ()] is a Hamiltonian Monte Carlo kernel.\n\n [integrator] defaults to {!leapfrog}. [num_leapfrog] defaults to [20]. *)\n\nval nuts_kernel :\n ?integrator:integrator ->\n ?max_depth:int ->\n step_size:float ->\n metric:metric ->\n unit ->\n kernel\n(** [nuts_kernel ~step_size ~metric ()] is a No-U-Turn Sampler kernel.\n\n NUTS automatically adapts the trajectory length using a binary tree\n expansion with U-turn detection. This eliminates the [num_leapfrog]\n parameter of {!hmc_kernel}.\n\n [integrator] defaults to {!leapfrog}. [max_depth] defaults to [10]. *)\n\n(** {1:sampling Sampling} *)\n\ntype stats = {\n accept_rate : float; (** Mean acceptance rate during sampling. *)\n step_size : float; (** Final adapted step size. *)\n num_divergent : int; (** Number of divergent transitions. *)\n}\n(** The type for aggregate sampling statistics. *)\n\ntype result = {\n samples : Nx.float64_t; (** Shape [[n; dim]]. *)\n log_densities : Nx.float64_t; (** Shape [[n]]. *)\n stats : stats;\n}\n(** The type for sampling results. *)\n\nval sample :\n ?step_size:float ->\n ?target_accept:float ->\n ?num_warmup:int ->\n ?report:(step:int -> state -> info -> unit) ->\n n:int ->\n (Nx.float64_t -> Nx.float64_t) ->\n Nx.float64_t ->\n (step_size:float -> metric:metric -> kernel) ->\n result\n(** [sample ~n log_prob init make_kernel] draws [n] samples from the\n distribution with unnormalized log-density [log_prob], starting at [init].\n\n During [num_warmup] iterations (discarded), step size and mass matrix are\n adapted using Stan-style window adaptation: an initial fast phase (step size\n only), doubling slow windows (step size + mass matrix with regularized\n Welford estimation), and a final fast phase.\n\n [step_size] defaults to [0.01]. [target_accept] defaults to [0.65].\n [num_warmup] defaults to [n / 2]. [report] is called after each step with\n negative step numbers during warmup. *)\n\nval hmc :\n ?step_size:float ->\n ?target_accept:float ->\n ?num_leapfrog:int ->\n ?num_warmup:int ->\n n:int ->\n (Nx.float64_t -> Nx.float64_t) ->\n Nx.float64_t ->\n result\n(** [hmc ~n log_prob init] draws [n] samples using Hamiltonian Monte Carlo with\n window adaptation.\n\n [step_size] defaults to [0.01]. [target_accept] defaults to [0.65].\n [num_leapfrog] defaults to [20]. [num_warmup] defaults to [n / 2]. *)\n\nval nuts :\n ?step_size:float ->\n ?target_accept:float ->\n ?max_depth:int ->\n ?num_warmup:int ->\n n:int ->\n (Nx.float64_t -> Nx.float64_t) ->\n Nx.float64_t ->\n result\n(** [nuts ~n log_prob init] draws [n] samples using the No-U-Turn Sampler with\n window adaptation.\n\n [step_size] defaults to [0.01]. [target_accept] defaults to [0.80].\n [max_depth] defaults to [10]. [num_warmup] defaults to [n / 2]. *)\n\n(** {1:diagnostics Diagnostics} *)\n\nval ess : Nx.float64_t -> Nx.float64_t\n(** [ess samples] is the effective sample size for each parameter. [samples] has\n shape [[n; dim]], returns shape [[dim]]. Computed via autocorrelation with\n the initial monotone sequence estimator. *)\n\nval rhat : Nx.float64_t array -> Nx.float64_t\n(** [rhat chains] is the split R-hat convergence diagnostic for each parameter.\n Each chain has shape [[n; dim]], returns shape [[dim]]. Values close to\n [1.0] indicate convergence; above [1.01] suggests the chains have not mixed.\n*)\n"
},
{
"path": "packages/norn/lib/nuts.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* No-U-Turn Sampler (Hoffman & Gelman 2014). Recursive binary tree expansion\n with multinomial sampling and generalised U-turn detection (Betancourt 2013).\n Follows BlackJAX's implementation strictly: - inner recursion:\n progressive_uniform_sampling (multinomial) - outer expansion:\n progressive_biased_sampling - U-turn: rho = m_sum - (m_right + m_left)/2,\n dot(v, rho) <= 0 *)\n\nlet f64 = Nx.float64\n\n(* Proposal: sampled state and acceptance statistics. *)\n\ntype proposal = {\n q : Nx.float64_t;\n lp : float;\n g : Nx.float64_t;\n energy : float;\n weight : float;\n sum_log_p_accept : float;\n}\n\n(* Trajectory: endpoints and accumulated momentum. *)\n\ntype trajectory = {\n left_q : Nx.float64_t;\n left_p : Nx.float64_t;\n left_g : Nx.float64_t;\n right_q : Nx.float64_t;\n right_p : Nx.float64_t;\n right_g : Nx.float64_t;\n momentum_sum : Nx.float64_t;\n num_states : int;\n}\n\nlet log_add_exp a b =\n if a = Float.neg_infinity then b\n else if b = Float.neg_infinity then a\n else if a >= b then a +. Float.log (1.0 +. Float.exp (b -. a))\n else b +. Float.log (1.0 +. Float.exp (a -. b))\n\n(* One leapfrog step → leaf proposal + trajectory. *)\nlet build_leaf (integrator : Internal.integrator) (metric : Internal.metric)\n direction step_size grad_log_prob q p g initial_energy =\n let eps = Float.of_int direction *. step_size in\n let q', p', lp', g' = integrator metric.scale q p g grad_log_prob eps in\n let energy = -.lp' +. metric.kinetic_energy p' in\n let delta = initial_energy -. energy in\n let delta = if Float.is_nan delta then Float.neg_infinity else delta in\n let proposal =\n {\n q = q';\n lp = lp';\n g = g';\n energy;\n weight = delta;\n sum_log_p_accept = Float.min delta 0.0;\n }\n in\n let trajectory =\n {\n left_q = q';\n left_p = p';\n left_g = g';\n right_q = q';\n right_p = p';\n right_g = g';\n momentum_sum = p';\n num_states = 1;\n }\n in\n let is_diverging = delta < -1000.0 in\n (proposal, trajectory, is_diverging)\n\n(* progressive_uniform_sampling: multinomial for inner tree building. *)\nlet uniform_sample prop new_prop =\n let p_accept = 1.0 /. (1.0 +. Float.exp (prop.weight -. new_prop.weight)) in\n let u = Nx.item [] (Nx.rand f64 [||]) in\n let weight = log_add_exp prop.weight new_prop.weight in\n let sum_log_p_accept =\n log_add_exp prop.sum_log_p_accept new_prop.sum_log_p_accept\n in\n if u < p_accept then { new_prop with weight; sum_log_p_accept }\n else { prop with weight; sum_log_p_accept }\n\n(* progressive_biased_sampling: for outer tree doubling. *)\nlet biased_sample prop new_prop =\n let p_accept = Float.min 1.0 (Float.exp (new_prop.weight -. prop.weight)) in\n let u = Nx.item [] (Nx.rand f64 [||]) in\n let weight = log_add_exp prop.weight new_prop.weight in\n let sum_log_p_accept =\n log_add_exp prop.sum_log_p_accept new_prop.sum_log_p_accept\n in\n if u < p_accept then { new_prop with weight; sum_log_p_accept }\n else { prop with weight; sum_log_p_accept }\n\nlet merge_trajectories direction traj new_traj =\n let l, r = if direction > 0 then (traj, new_traj) else (new_traj, traj) in\n {\n left_q = l.left_q;\n left_p = l.left_p;\n left_g = l.left_g;\n right_q = r.right_q;\n right_p = r.right_p;\n right_g = r.right_g;\n momentum_sum = Nx.add l.momentum_sum r.momentum_sum;\n num_states = l.num_states + r.num_states;\n }\n\n(* Recursive tree building — buildtree_integrate from BlackJAX\n trajectory.py:dynamic_recursive_integration. *)\nlet rec build_tree integrator metric step_size grad_log_prob q p g depth\n direction initial_energy =\n if depth = 0 then\n let prop, traj, is_div =\n build_leaf integrator metric direction step_size grad_log_prob q p g\n initial_energy\n in\n (prop, traj, is_div, false)\n else\n let half = depth - 1 in\n let prop, traj, is_div, is_turn =\n build_tree integrator metric step_size grad_log_prob q p g half direction\n initial_energy\n in\n if is_div || is_turn then (prop, traj, is_div, is_turn)\n else\n let q', p', g' =\n if direction > 0 then (traj.right_q, traj.right_p, traj.right_g)\n else (traj.left_q, traj.left_p, traj.left_g)\n in\n let new_prop, new_traj, new_div, new_turn =\n build_tree integrator metric step_size grad_log_prob q' p' g' half\n direction initial_energy\n in\n let merged = merge_trajectories direction traj new_traj in\n if new_turn then\n (* Second half turning: keep old proposal state/weight, accumulate\n sum_log_p_accept for consistent acceptance_rate. *)\n let slpa =\n log_add_exp prop.sum_log_p_accept new_prop.sum_log_p_accept\n in\n ({ prop with sum_log_p_accept = slpa }, merged, new_div, true)\n else\n (* Check U-turn on merged trajectory *)\n let turning =\n metric.is_turning merged.left_p merged.right_p merged.momentum_sum\n in\n (* Always sample when second half is not turning *)\n let sampled = uniform_sample prop new_prop in\n (sampled, merged, new_div, turning)\n\n(* Outer expansion loop — dynamic_multiplicative_expansion from BlackJAX\n trajectory.py. *)\nlet step (integrator : Internal.integrator) (metric : Internal.metric) step_size\n max_depth (state : Internal.state) log_density_fn =\n let grad_log_prob q =\n let lp, g = Rune.value_and_grad log_density_fn q in\n (Nx.item [] lp, g)\n in\n let dim = Nx.numel state.position in\n let p0 = metric.sample_momentum dim in\n let ke0 = metric.kinetic_energy p0 in\n let initial_energy = -.state.log_density +. ke0 in\n let proposal =\n ref\n {\n q = state.position;\n lp = state.log_density;\n g = state.grad_log_density;\n energy = initial_energy;\n weight = 0.0;\n sum_log_p_accept = Float.neg_infinity;\n }\n in\n let trajectory =\n ref\n {\n left_q = state.position;\n left_p = p0;\n left_g = state.grad_log_density;\n right_q = state.position;\n right_p = p0;\n right_g = state.grad_log_density;\n momentum_sum = p0;\n num_states = 0;\n }\n in\n let depth = ref 0 in\n let diverging = ref false in\n let turning = ref false in\n while !depth < max_depth && (not !diverging) && not !turning do\n let direction = if Nx.item [] (Nx.rand f64 [||]) < 0.5 then -1 else 1 in\n let q, p, g =\n if direction > 0 then\n (!trajectory.right_q, !trajectory.right_p, !trajectory.right_g)\n else (!trajectory.left_q, !trajectory.left_p, !trajectory.left_g)\n in\n let sub_prop, sub_traj, sub_div, sub_turn =\n build_tree integrator metric step_size grad_log_prob q p g !depth\n direction initial_energy\n in\n (* Update proposal: biased sampling unless subtree diverged or turned *)\n if sub_div || sub_turn then\n proposal :=\n {\n !proposal with\n sum_log_p_accept =\n log_add_exp !proposal.sum_log_p_accept sub_prop.sum_log_p_accept;\n }\n else proposal := biased_sample !proposal sub_prop;\n (* Always merge trajectory *)\n trajectory := merge_trajectories direction !trajectory sub_traj;\n (* Check U-turn on full trajectory *)\n let full_turn =\n metric.is_turning !trajectory.left_p !trajectory.right_p\n !trajectory.momentum_sum\n in\n diverging := sub_div;\n turning := sub_turn || full_turn;\n incr depth\n done;\n let p = !proposal in\n let t = !trajectory in\n let new_state : Internal.state =\n { position = p.q; log_density = p.lp; grad_log_density = p.g }\n in\n let n_states = Float.of_int (max 1 t.num_states) in\n let acceptance_rate = Float.exp p.sum_log_p_accept /. n_states in\n let info : Internal.info =\n {\n acceptance_rate;\n is_divergent = !diverging;\n energy = initial_energy;\n num_integration_steps = t.num_states;\n }\n in\n (new_state, info)\n"
},
{
"path": "packages/norn/test/debug_nuts.ml",
"content": "let f64 = Nx.float64\nlet true_mean = Nx.create f64 [| 2 |] [| 3.0; -1.0 |]\nlet true_cov = Nx.create f64 [| 2; 2 |] [| 1.0; 0.5; 0.5; 2.0 |]\nlet cov_inv = Nx.inv true_cov\n\nlet log_prob x =\n let d = Nx.sub x true_mean in\n let dt = Nx.reshape [| 1; 2 |] d in\n let mahal = Nx.matmul (Nx.matmul dt cov_inv) (Nx.reshape [| 2; 1 |] d) in\n Nx.mul_s (Nx.reshape [||] mahal) (-0.5)\n\nlet compute_stats name positions n =\n let nf = Float.of_int n in\n let sum0 = ref 0.0 in\n let sum1 = ref 0.0 in\n for i = 0 to n - 1 do\n let x0, x1 = positions.(i) in\n sum0 := !sum0 +. x0;\n sum1 := !sum1 +. x1\n done;\n let mean0 = !sum0 /. nf in\n let mean1 = !sum1 /. nf in\n let var0 = ref 0.0 in\n let var1 = ref 0.0 in\n for i = 0 to n - 1 do\n let x0, x1 = positions.(i) in\n var0 := !var0 +. ((x0 -. mean0) *. (x0 -. mean0));\n var1 := !var1 +. ((x1 -. mean1) *. (x1 -. mean1))\n done;\n let var0 = !var0 /. nf in\n let var1 = !var1 /. nf in\n Printf.printf \"%s:\\n\" name;\n Printf.printf \" mean = [%.4f, %.4f] (true: [3.0, -1.0])\\n\" mean0 mean1;\n Printf.printf \" var = [%.4f, %.4f] (true: [1.0, 2.0])\\n\" var0 var1;\n Printf.printf \" var error = [%.1f%%, %.1f%%]\\n\"\n (100.0 *. Float.abs (var0 -. 1.0) /. 1.0)\n (100.0 *. Float.abs (var1 -. 2.0) /. 2.0)\n\nlet () =\n Printf.printf \"=== Comparison: Norn vs BlackJAX ===\\n\\n\";\n Printf.printf \"Target: 2D Gaussian, mean=[3,-1], cov=[[1,0.5],[0.5,2]]\\n\\n\";\n\n (* --- HMC: match BlackJAX: 1000 warmup + 3000 samples, fixed --- *)\n Nx.Rng.run ~seed:42 (fun () ->\n let metric = Norn.unit_metric 2 in\n let kernel = Norn.hmc_kernel ~step_size:0.1 ~num_leapfrog:20 ~metric () in\n let init = Nx.zeros f64 [| 2 |] in\n let state = ref (kernel.init init log_prob) in\n for _ = 1 to 1000 do\n let s, _ = kernel.step !state log_prob in\n state := s\n done;\n let positions = Array.make 3000 (0.0, 0.0) in\n for i = 0 to 2999 do\n let s, _ = kernel.step !state log_prob in\n state := s;\n positions.(i) <- (Nx.item [ 0 ] s.position, Nx.item [ 1 ] s.position)\n done;\n compute_stats \"Norn HMC (fixed eps=0.1, 1000w+3000s)\" positions 3000;\n Printf.printf\n \" BlackJAX: mean=[3.0048, -0.9621] var=[0.9896, 2.0819]\\n\\n\");\n\n (* --- NUTS: match BlackJAX: 1000 warmup + 3000 samples, fixed --- *)\n Nx.Rng.run ~seed:42 (fun () ->\n let metric = Norn.unit_metric 2 in\n let kernel = Norn.nuts_kernel ~step_size:0.1 ~metric () in\n let init = Nx.zeros f64 [| 2 |] in\n let state = ref (kernel.init init log_prob) in\n for _ = 1 to 1000 do\n let s, _ = kernel.step !state log_prob in\n state := s\n done;\n let positions = Array.make 3000 (0.0, 0.0) in\n let total_lf = ref 0 in\n for i = 0 to 2999 do\n let s, info = kernel.step !state log_prob in\n state := s;\n total_lf := !total_lf + info.num_integration_steps;\n positions.(i) <- (Nx.item [ 0 ] s.position, Nx.item [ 1 ] s.position)\n done;\n compute_stats \"Norn NUTS (fixed eps=0.1, 1000w+3000s)\" positions 3000;\n Printf.printf \" avg leapfrog/step = %.1f\\n\"\n (Float.of_int !total_lf /. 3000.0);\n Printf.printf\n \" BlackJAX: mean=[3.0297, -0.9192] var=[1.0453, 2.1354]\\n\\n\");\n\n (* --- NUTS adapted: warmup trace + 100 samples --- *)\n Nx.Rng.run ~seed:42 (fun () ->\n let report ~step (_state : Norn.state) (info : Norn.info) =\n if step < 0 then begin\n let ws = 1000 + step + 1 in\n if ws <= 10 || ws mod 100 = 0 then\n Printf.printf \" warmup %4d: accept=%.4f lf=%d\\n\" ws\n info.acceptance_rate info.num_integration_steps\n end\n in\n let result =\n Norn.sample ~step_size:0.1 ~target_accept:0.80 ~num_warmup:1000 ~report\n ~n:100 log_prob (Nx.zeros f64 [| 2 |]) (fun ~step_size ~metric ->\n Norn.nuts_kernel ~step_size ~metric ())\n in\n Printf.printf \"\\nAdapted: step_size=%.6f accept=%.4f divergent=%d\\n\"\n result.stats.step_size result.stats.accept_rate\n result.stats.num_divergent)\n"
},
{
"path": "packages/norn/test/dune",
"content": "(test\n (name test_norn)\n (package norn)\n (libraries nx rune norn windtrap))\n\n(executable\n (name debug_nuts)\n (libraries nx rune norn))\n"
},
{
"path": "packages/norn/test/test_blackjax_ref.py",
"content": "\"\"\"Reference test: 2D correlated Gaussian with BlackJAX NUTS and HMC.\"\"\"\nimport jax\nimport jax.numpy as jnp\nimport blackjax\n\ntrue_mean = jnp.array([3.0, -1.0])\ntrue_cov = jnp.array([[1.0, 0.5], [0.5, 2.0]])\ncov_inv = jnp.linalg.inv(true_cov)\n\ndef log_prob(x):\n d = x - true_mean\n return -0.5 * d @ cov_inv @ d\n\n# --- HMC ---\nkey = jax.random.key(42)\ninv_mass = jnp.ones(2)\nhmc_alg = blackjax.hmc(log_prob, step_size=0.1, inverse_mass_matrix=inv_mass, num_integration_steps=20)\nstate = hmc_alg.init(jnp.zeros(2))\nstep_fn = jax.jit(hmc_alg.step)\n\nsamples_hmc = []\nfor i in range(4000): # 1000 warmup + 3000 samples\n key = jax.random.fold_in(key, i)\n state, info = step_fn(key, state)\n if i >= 1000:\n samples_hmc.append(state.position)\n\nsamples_hmc = jnp.stack(samples_hmc)\nprint(f\"HMC (no adaptation, fixed step_size=0.1):\")\nprint(f\" mean = {jnp.mean(samples_hmc, axis=0)}\")\nprint(f\" var = {jnp.var(samples_hmc, axis=0)}\")\nprint()\n\n# --- NUTS ---\nkey = jax.random.key(42)\nnuts_alg = blackjax.nuts(log_prob, step_size=0.1, inverse_mass_matrix=inv_mass)\nstate = nuts_alg.init(jnp.zeros(2))\nstep_fn = jax.jit(nuts_alg.step)\n\nsamples_nuts = []\nfor i in range(4000):\n key = jax.random.fold_in(key, i)\n state, info = step_fn(key, state)\n if i >= 1000:\n samples_nuts.append(state.position)\n\nsamples_nuts = jnp.stack(samples_nuts)\nprint(f\"NUTS (no adaptation, fixed step_size=0.1):\")\nprint(f\" mean = {jnp.mean(samples_nuts, axis=0)}\")\nprint(f\" var = {jnp.var(samples_nuts, axis=0)}\")\nprint()\n\n# --- NUTS with window adaptation ---\nkey = jax.random.key(42)\nwarmup = blackjax.window_adaptation(blackjax.nuts, log_prob, num_steps=1000)\nkey, warmup_key = jax.random.split(key)\n(adapted_state, adapted_params), _ = warmup.run(warmup_key, jnp.zeros(2))\nprint(f\"Window adaptation results:\")\nprint(f\" step_size = {adapted_params['step_size']}\")\nprint(f\" inv_mass = {adapted_params['inverse_mass_matrix']}\")\n\nnuts_adapted = blackjax.nuts(log_prob, **adapted_params)\nstep_fn = jax.jit(nuts_adapted.step)\nstate = adapted_state\n\nsamples_adapted = []\nfor i in range(3000):\n key = jax.random.fold_in(key, i)\n state, info = step_fn(key, state)\n samples_adapted.append(state.position)\n\nsamples_adapted = jnp.stack(samples_adapted)\nprint(f\"NUTS (with window adaptation):\")\nprint(f\" mean = {jnp.mean(samples_adapted, axis=0)}\")\nprint(f\" var = {jnp.var(samples_adapted, axis=0)}\")\n"
},
{
"path": "packages/norn/test/test_norn.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Windtrap\n\nlet f64 = Nx.float64\n\n(* 2D correlated Gaussian: mean [3, -1], covariance [[1, 0.5], [0.5, 2]] *)\nlet true_mean = Nx.create f64 [| 2 |] [| 3.0; -1.0 |]\nlet true_cov = Nx.create f64 [| 2; 2 |] [| 1.0; 0.5; 0.5; 2.0 |]\nlet cov_inv = Nx.inv true_cov\n\nlet log_prob x =\n let d = Nx.sub x true_mean in\n let dt = Nx.reshape [| 1; 2 |] d in\n let mahal = Nx.matmul (Nx.matmul dt cov_inv) (Nx.reshape [| 2; 1 |] d) in\n Nx.mul_s (Nx.reshape [||] mahal) (-0.5)\n\nlet check_result msg result =\n is_true\n ~msg:\n (Printf.sprintf \"%s: accept rate %.2f > 0.4\" msg\n result.Norn.stats.accept_rate)\n (result.stats.accept_rate > 0.4);\n let sample_mean = Nx.mean ~axes:[ 0 ] result.samples in\n for i = 0 to 1 do\n let sm = Nx.item [ i ] sample_mean in\n let tm = Nx.item [ i ] true_mean in\n is_true\n ~msg:(Printf.sprintf \"%s: mean[%d]: %.2f ~ %.2f\" msg i sm tm)\n (Float.abs (sm -. tm) < 0.5)\n done;\n let centered = Nx.sub result.samples sample_mean in\n let n = Float.of_int ((Nx.shape result.samples).(0) - 1) in\n let sample_cov =\n Nx.div_s (Nx.matmul (Nx.matrix_transpose centered) centered) n\n in\n for i = 0 to 1 do\n let sc = Nx.item [ i; i ] sample_cov in\n let tc = Nx.item [ i; i ] true_cov in\n is_true\n ~msg:(Printf.sprintf \"%s: var[%d]: %.2f ~ %.2f (within 60%%)\" msg i sc tc)\n (Float.abs (sc -. tc) /. tc < 0.6)\n done\n\nlet test_hmc () =\n Nx.Rng.run ~seed:42 (fun () ->\n let init = Nx.zeros f64 [| 2 |] in\n let result =\n Norn.hmc ~step_size:0.1 ~num_leapfrog:20 ~num_warmup:200 ~n:500 log_prob\n init\n in\n check_result \"HMC\" result)\n\nlet test_nuts () =\n Nx.Rng.run ~seed:42 (fun () ->\n let init = Nx.zeros f64 [| 2 |] in\n let result =\n Norn.nuts ~step_size:0.5 ~max_depth:6 ~num_warmup:500 ~n:800 log_prob\n init\n in\n check_result \"NUTS\" result)\n\nlet test_kernel_api () =\n Nx.Rng.run ~seed:42 (fun () ->\n let init = Nx.zeros f64 [| 2 |] in\n let metric = Norn.unit_metric 2 in\n let kernel = Norn.hmc_kernel ~step_size:0.1 ~metric () in\n let state = kernel.init init log_prob in\n is_true ~msg:\"init log_density is finite\"\n (Float.is_finite state.log_density);\n let state', info = kernel.step state log_prob in\n is_true ~msg:\"step produces finite log_density\"\n (Float.is_finite state'.log_density);\n is_true ~msg:\"acceptance_rate in [0, 1]\"\n (info.acceptance_rate >= 0.0 && info.acceptance_rate <= 1.0))\n\nlet test_sample_with_kernel () =\n Nx.Rng.run ~seed:42 (fun () ->\n let init = Nx.zeros f64 [| 2 |] in\n let result =\n Norn.sample ~step_size:0.1 ~num_warmup:200 ~n:500 log_prob init\n (fun ~step_size ~metric -> Norn.hmc_kernel ~step_size ~metric ())\n in\n check_result \"sample+kernel\" result)\n\nlet test_diagnostics () =\n Nx.Rng.run ~seed:42 (fun () ->\n let init = Nx.zeros f64 [| 2 |] in\n let chains =\n Array.init 4 (fun i ->\n Nx.Rng.run ~seed:i (fun () ->\n Norn.nuts ~step_size:0.1 ~num_warmup:500 ~n:1000 log_prob init))\n in\n let chain_samples = Array.map (fun r -> r.Norn.samples) chains in\n let r = Norn.rhat chain_samples in\n for d = 0 to 1 do\n let rv = Nx.item [ d ] r in\n is_true ~msg:(Printf.sprintf \"rhat[%d]: %.3f < 1.1\" d rv) (rv < 1.1)\n done;\n let e = Norn.ess chain_samples.(0) in\n for d = 0 to 1 do\n let ev = Nx.item [ d ] e in\n is_true ~msg:(Printf.sprintf \"ess[%d]: %.0f > 50\" d ev) (ev > 50.0)\n done)\n\nlet () =\n run \"Norn\"\n [\n test \"HMC: 2D Gaussian\" test_hmc;\n test \"NUTS: 2D Gaussian\" test_nuts;\n test \"Kernel API\" test_kernel_api;\n test \"Sample with kernel\" test_sample_with_kernel;\n test \"Diagnostics\" test_diagnostics;\n ]\n"
},
{
"path": "packages/nx/README.md",
"content": "# Nx\n\nN-dimensional array library for OCaml.\n\nNx is the core component of the Raven ecosystem, providing efficient\nnumerical computation with multi-device support. It offers NumPy-like\nfunctionality with the benefits of OCaml's strong static type system.\n\n## Features\n\n- Multi-dimensional arrays (tensors) with arbitrary rank\n- Support for data types: float16, float32, float64, int8, int16, int32,\n int64, uint8, uint16, complex32, complex64\n- Flexible memory layouts: C-contiguous and strided\n- Zero-copy slicing, reshaping, and broadcasting\n- Element-wise and scalar operations (add, sub, mul, div, map, etc.)\n- Linear algebra routines (`dot`, matrix multiplication, transpose,\n sum, mean, argmax, etc.)\n- Optimized CPU backend; pure OCaml interface leveraging Bigarray\n- I/O support: image formats (PNG, JPEG), NumPy files (.npy, .npz)\n- Seamless integration with the Raven ecosystem: `sowilo`, `quill`, `hugin`, etc.\n\n## Quick Start\n\n```ocaml\nopen Nx\n\n(* Create a 2x3 tensor *)\nlet a = create float32 [|2;3|] [|1.; 2.; 3.; 4.; 5.; 6.|]\n\n(* Fill a tensor with ones *)\nlet b = full float32 [|2;3|] 1.0\n\n(* Element-wise addition *)\nlet c = add a b\n\n(* Matrix multiplication *)\nlet x = create float32 [|2;3|] [|1.;2.;3.;4.;5.;6.|]\nlet y = create float32 [|3;2|] [|7.;8.;9.;10.;11.;12.|]\nlet z = dot x y\n\n(* Reduction: sum across an axis *)\nlet s = sum ~axes:[|1|] x\n```\n\n## Contributing\n\nSee the [Raven monorepo README](../README.md) for contribution guidelines.\n\n## License\n\nISC License. See [LICENSE](../LICENSE) for details.\n"
},
{
"path": "packages/nx/bench/README.md",
"content": "# Nx Benchmarks\n\nThis directory contains benchmarks for the nx library. We provide comparative benchmarks against NumPy.\n\nAdditional focused suites live under subdirectories:\n\n- `conv2d/` for convolutional workloads\n- `einsum/` for tensor contractions\n- `matmul/` for dense matrix multiplication\n\n## Results Nx\n\n```\n┌────────────────────────────┬──────────┬──────────┬─────────┬─────────┬────────────┐\n│ Name │ Wall/Run │ CPU/Run │ mWd/Run │ Speedup │ vs Fastest │\n├────────────────────────────┼──────────┼──────────┼─────────┼─────────┼────────────┤\n│ Transpose 100x100 f32 (Nx) │ 132.92ns │ 174.70ns │ 114.00w │ 1.00x │ 100% │\n│ Transpose 200x200 f64 (Nx) │ 144.19ns │ 162.57ns │ 114.00w │ 0.92x │ 108% │\n│ Transpose 50x50 f32 (Nx) │ 145.52ns │ 184.50ns │ 114.00w │ 0.91x │ 109% │\n│ Transpose 500x500 f64 (Nx) │ 147.02ns │ 166.86ns │ 114.00w │ 0.90x │ 111% │\n│ Transpose 100x100 f64 (Nx) │ 147.10ns │ 232.42ns │ 114.00w │ 0.90x │ 111% │\n│ Transpose 500x500 f32 (Nx) │ 148.00ns │ 162.61ns │ 114.00w │ 0.90x │ 111% │\n│ Transpose 200x200 f32 (Nx) │ 150.68ns │ 170.96ns │ 114.00w │ 0.88x │ 113% │\n│ Transpose 50x50 f64 (Nx) │ 158.90ns │ 161.09ns │ 114.00w │ 0.84x │ 120% │\n│ Mul 50x50 f32 (Nx) │ 1.61μs │ 1.61μs │ 808.00w │ 0.08x │ 1211% │\n│ Mul 50x50 f64 (Nx) │ 1.74μs │ 1.68μs │ 808.00w │ 0.08x │ 1311% │\n│ Add 50x50 f64 (Nx) │ 1.77μs │ 1.82μs │ 808.00w │ 0.07x │ 1335% │\n│ Add 50x50 f32 (Nx) │ 1.83μs │ 1.83μs │ 808.00w │ 0.07x │ 1374% │\n│ Mul 100x100 f32 (Nx) │ 2.41μs │ 2.52μs │ 808.00w │ 0.06x │ 1814% │\n│ Add 100x100 f32 (Nx) │ 2.69μs │ 2.69μs │ 808.00w │ 0.05x │ 2022% │\n│ Sum 50x50 f64 (Nx) │ 3.23μs │ 3.25μs │ 416.00w │ 0.04x │ 2432% │\n│ Sum 50x50 f32 (Nx) │ 3.27μs │ 3.24μs │ 416.00w │ 0.04x │ 2457% │\n│ Mul 100x100 f64 (Nx) │ 22.87μs │ 22.72μs │ 801.00w │ 0.01x │ 17206% │\n│ Add 100x100 f64 (Nx) │ 24.60μs │ 24.55μs │ 801.00w │ 0.01x │ 18510% │\n│ Sum 100x100 f64 (Nx) │ 87.16μs │ 178.36μs │ 416.00w │ 0.00x │ 65571% │\n│ Sum 100x100 f32 (Nx) │ 87.68μs │ 178.89μs │ 416.00w │ 0.00x │ 65966% │\n│ Sum 200x200 f64 (Nx) │ 97.11μs │ 220.67μs │ 416.00w │ 0.00x │ 73060% │\n│ Sum 200x200 f32 (Nx) │ 98.78μs │ 218.97μs │ 416.00w │ 0.00x │ 74315% │\n│ Mul 200x200 f32 (Nx) │ 105.68μs │ 152.89μs │ 801.00w │ 0.00x │ 79507% │\n│ Add 200x200 f32 (Nx) │ 109.71μs │ 161.79μs │ 801.00w │ 0.00x │ 82541% │\n│ Sum 500x500 f32 (Nx) │ 135.81μs │ 425.91μs │ 416.00w │ 0.00x │ 102173% │\n│ Sum 500x500 f64 (Nx) │ 139.32μs │ 427.21μs │ 416.00w │ 0.00x │ 104814% │\n│ Mul 200x200 f64 (Nx) │ 157.28μs │ 210.49μs │ 801.00w │ 0.00x │ 118332% │\n│ Add 200x200 f64 (Nx) │ 162.52μs │ 223.83μs │ 801.00w │ 0.00x │ 122270% │\n│ Mul 500x500 f32 (Nx) │ 195.66μs │ 271.61μs │ 801.00w │ 0.00x │ 147200% │\n│ Add 500x500 f32 (Nx) │ 204.88μs │ 301.06μs │ 801.00w │ 0.00x │ 154141% │\n│ Mul 500x500 f64 (Nx) │ 277.63μs │ 534.55μs │ 801.00w │ 0.00x │ 208873% │\n│ Add 500x500 f64 (Nx) │ 284.96μs │ 571.38μs │ 801.00w │ 0.00x │ 214386% │\n└────────────────────────────┴──────────┴──────────┴─────────┴─────────┴────────────┘\n```\n\n## Results NumPy\n\n```\n┌───────────────────────────────┬──────────┬──────────┬─────────┬─────────┬────────────┐\n│ Name │ Wall/Run │ CPU/Run │ mWd/Run │ Speedup │ vs Fastest │\n├───────────────────────────────┼──────────┼──────────┼─────────┼─────────┼────────────┤\n│ Transpose 200x200 f64 (NumPy) │ 309.12ns │ 308.61ns │ 0.04w │ 1.00x │ 100% │\n│ Transpose 500x500 f64 (NumPy) │ 310.86ns │ 310.27ns │ 0.06w │ 0.99x │ 101% │\n│ Transpose 100x100 f64 (NumPy) │ 311.56ns │ 309.53ns │ 0.04w │ 0.99x │ 101% │\n│ Transpose 50x50 f32 (NumPy) │ 311.76ns │ 311.08ns │ 0.04w │ 0.99x │ 101% │\n│ Transpose 50x50 f64 (NumPy) │ 313.21ns │ 312.09ns │ 0.04w │ 0.99x │ 101% │\n│ Transpose 200x200 f32 (NumPy) │ 313.63ns │ 312.67ns │ 0.04w │ 0.99x │ 101% │\n│ Transpose 500x500 f32 (NumPy) │ 315.64ns │ 313.60ns │ 0.05w │ 0.98x │ 102% │\n│ Transpose 100x100 f32 (NumPy) │ 322.03ns │ 315.85ns │ 0.04w │ 0.96x │ 104% │\n│ Add 50x50 f32 (NumPy) │ 541.58ns │ 540.36ns │ 0.06w │ 0.57x │ 175% │\n│ Mul 50x50 f32 (NumPy) │ 547.03ns │ 539.71ns │ 0.06w │ 0.57x │ 177% │\n│ Add 50x50 f64 (NumPy) │ 688.90ns │ 688.46ns │ 0.08w │ 0.45x │ 223% │\n│ Mul 50x50 f64 (NumPy) │ 775.58ns │ 711.24ns │ 0.08w │ 0.40x │ 251% │\n│ Mul 100x100 f32 (NumPy) │ 1.06µs │ 1.06µs │ 0.11w │ 0.29x │ 342% │\n│ Add 100x100 f32 (NumPy) │ 1.08µs │ 1.05µs │ 0.12w │ 0.29x │ 351% │\n│ Sum 50x50 f64 (NumPy) │ 1.88µs │ 1.88µs │ 0.20w │ 0.16x │ 608% │\n│ Sum 50x50 f32 (NumPy) │ 1.90µs │ 1.89µs │ 0.20w │ 0.16x │ 615% │\n│ Mul 100x100 f64 (NumPy) │ 2.54µs │ 2.51µs │ 0.32w │ 0.12x │ 822% │\n│ Add 100x100 f64 (NumPy) │ 2.58µs │ 2.52µs │ 0.32w │ 0.12x │ 835% │\n│ Sum 100x100 f32 (NumPy) │ 3.16µs │ 3.16µs │ 0.34w │ 0.10x │ 1023% │\n│ Sum 100x100 f64 (NumPy) │ 3.16µs │ 3.16µs │ 0.34w │ 0.10x │ 1023% │\n│ Mul 200x200 f32 (NumPy) │ 4.39µs │ 4.37µs │ 0.55w │ 0.07x │ 1419% │\n│ Add 200x200 f32 (NumPy) │ 4.39µs │ 4.38µs │ 0.54w │ 0.07x │ 1421% │\n│ Add 200x200 f64 (NumPy) │ 8.30µs │ 8.21µs │ 0.95w │ 0.04x │ 2685% │\n│ Mul 200x200 f64 (NumPy) │ 8.45µs │ 8.32µs │ 0.97w │ 0.04x │ 2734% │\n│ Sum 200x200 f32 (NumPy) │ 9.06µs │ 8.94µs │ 0.92w │ 0.03x │ 2931% │\n│ Sum 200x200 f64 (NumPy) │ 9.53µs │ 9.50µs │ 1.26w │ 0.03x │ 3082% │\n│ Add 500x500 f32 (NumPy) │ 24.95µs │ 24.89µs │ 3.01w │ 0.01x │ 8071% │\n│ Mul 500x500 f32 (NumPy) │ 25.01µs │ 24.95µs │ 3.08w │ 0.01x │ 8090% │\n│ Sum 500x500 f32 (NumPy) │ 43.39µs │ 43.00µs │ 5.57w │ 0.01x │ 14036% │\n│ Sum 500x500 f64 (NumPy) │ 48.26µs │ 48.03µs │ 8.13w │ 0.01x │ 15614% │\n│ Add 500x500 f64 (NumPy) │ 159.37µs │ 159.03µs │ 23.29w │ 0.00x │ 51557% │\n│ Mul 500x500 f64 (NumPy) │ 167.38µs │ 166.97µs │ 28.98w │ 0.00x │ 54149% │\n└───────────────────────────────┴──────────┴──────────┴─────────┴─────────┴────────────┘\n```\n"
},
{
"path": "packages/nx/bench/bench_numpy.py",
"content": "from __future__ import annotations\n\nimport sys\nfrom pathlib import Path\nfrom typing import Any, Callable, Iterable, List, Sequence, Tuple\n\nimport numpy as np\n\n_SCRIPTS_DIR = Path(__file__).resolve().parent\nwhile not (_SCRIPTS_DIR / \"dune-project\").exists():\n _SCRIPTS_DIR = _SCRIPTS_DIR.parent\n_SCRIPTS_DIR = _SCRIPTS_DIR / \"scripts\"\nif str(_SCRIPTS_DIR) not in sys.path:\n sys.path.insert(0, str(_SCRIPTS_DIR))\n\nimport ubench # type: ignore\n\n\nSIZES: Sequence[int] = (50, 100, 200, 500)\nDTYPES: Sequence[np.dtype] = (np.float32, np.float64)\nBACKEND_NAME = \"NumPy\"\n_RNG = np.random.default_rng(seed=0)\n\n\ndef _dtype_label(dtype: np.dtype) -> str:\n if dtype == np.float32:\n return \"f32\"\n if dtype == np.float64:\n return \"f64\"\n return str(dtype)\n\n\ndef _benchmark_name(op_name: str, size: int, dtype: np.dtype) -> str:\n return f\"{op_name} {size}x{size} {_dtype_label(dtype)} ({BACKEND_NAME})\"\n\n\ndef _numpy_operations(\n size: int, dtype: np.dtype\n) -> Iterable[Tuple[str, Callable[[], None]]]:\n a = _RNG.random((size, size), dtype=dtype)\n b = _RNG.random((size, size), dtype=dtype)\n\n ops: List[Tuple[str, Callable[[], None]]] = [\n (\"Add\", lambda a=a, b=b: np.add(a, b)),\n (\"Mul\", lambda a=a, b=b: np.multiply(a, b)),\n ]\n\n ops.extend(\n [\n (\"Sum\", lambda a=a: np.sum(a)),\n (\"Transpose\", lambda a=a: np.transpose(a)),\n ]\n )\n return ops\n\n\ndef build_benchmarks() -> List[Any]:\n benchmarks: List[Any] = []\n for size in SIZES:\n for dtype in DTYPES:\n for op_name, fn in _numpy_operations(size, dtype):\n bench_name = _benchmark_name(op_name, size, dtype)\n benchmarks.append(ubench.bench(bench_name, fn))\n return benchmarks\n\n\ndef default_config() -> ubench.Config:\n return (\n ubench.Config.default()\n .time_limit(1.0)\n .warmup(1)\n .min_measurements(5)\n .min_cpu(0.01)\n .geometric_scale(1.3)\n .gc_stabilization(False)\n .build()\n )\n\n\ndef main() -> None:\n benchmarks = build_benchmarks()\n # Mirror the OCaml defaults for fair comparisons with Nx benchmarks.\n config = default_config()\n ubench.run(benchmarks, config=config, output_format=\"pretty\", verbose=False)\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "packages/nx/bench/bench_nx.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Configuration *)\nlet sizes = [ 50; 100; 200; 500 ]\nlet backend_name = \"Nx\"\n\nlet benchmark_name op_name size dtype_label =\n Printf.sprintf \"%s %dx%d %s (%s)\" op_name size size dtype_label backend_name\n\nlet nx_operations_f32 ~size =\n let shape = [| size; size |] in\n let a = Nx.rand Nx.Float32 shape in\n let b = Nx.rand Nx.Float32 shape in\n\n let ops = [ (\"Add\", fun () -> Nx.add a b); (\"Mul\", fun () -> Nx.mul a b) ] in\n\n let ops =\n ops\n @ [ (\"Sum\", fun () -> Nx.sum a); (\"Transpose\", fun () -> Nx.transpose a) ]\n in\n\n ops\n\nlet nx_operations_f64 ~size =\n let shape = [| size; size |] in\n let a = Nx.rand Nx.Float64 shape in\n let b = Nx.rand Nx.Float64 shape in\n\n let ops = [ (\"Add\", fun () -> Nx.add a b); (\"Mul\", fun () -> Nx.mul a b) ] in\n\n let ops =\n ops\n @ [ (\"Sum\", fun () -> Nx.sum a); (\"Transpose\", fun () -> Nx.transpose a) ]\n in\n\n ops\n\nlet build_benchmarks () =\n let f32_benches = ref [] in\n let f64_benches = ref [] in\n List.iter\n (fun size ->\n let ops_f32 = nx_operations_f32 ~size in\n List.iter\n (fun (op_name, fn) ->\n let bench_name = benchmark_name op_name size \"f32\" in\n f32_benches := Thumper.bench bench_name fn :: !f32_benches)\n ops_f32;\n\n let ops_f64 = nx_operations_f64 ~size in\n List.iter\n (fun (op_name, fn) ->\n let bench_name = benchmark_name op_name size \"f64\" in\n f64_benches := Thumper.bench bench_name fn :: !f64_benches)\n ops_f64)\n sizes;\n [\n Thumper.group \"f32\" (List.rev !f32_benches);\n Thumper.group \"f64\" (List.rev !f64_benches);\n ]\n\nlet () =\n let benchmarks = build_benchmarks () in\n Thumper.run \"nx\" benchmarks\n"
},
{
"path": "packages/nx/bench/conv2d/README.md",
"content": "# Nx Conv2d Benchmarks\n\nComparative benchmarks of Nx conv2d operations against PyTorch.\n\n## Results Nx\n\n```\n┌───────────────────────────────────────┬──────────┬──────────┬─────────┬─────────┬────────────┐\n│ Name │ Wall/Run │ CPU/Run │ mWd/Run │ Speedup │ vs Fastest │\n├───────────────────────────────────────┼──────────┼──────────┼─────────┼─────────┼────────────┤\n│ Conv2d B1 C3->32 64x64 K3 f32 (Nx) │ 251.48μs │ 250.62μs │ 2.79kw │ 1.00x │ 100% │\n│ Conv2d B1 C3->32 64x64 K3 f64 (Nx) │ 359.48μs │ 357.43μs │ 2.79kw │ 0.70x │ 143% │\n│ Conv2d B8 C32->64 32x32 K3 f32 (Nx) │ 4.30ms │ 4.28ms │ 2.78kw │ 0.06x │ 1710% │\n│ Conv2d B16 C64->128 16x16 K3 f32 (Nx) │ 4.45ms │ 4.35ms │ 2.78kw │ 0.06x │ 1768% │\n│ Conv2d B8 C32->64 32x32 K3 f64 (Nx) │ 5.31ms │ 5.25ms │ 2.78kw │ 0.05x │ 2111% │\n│ Conv2d B16 C64->128 16x16 K3 f64 (Nx) │ 5.73ms │ 5.51ms │ 2.78kw │ 0.04x │ 2279% │\n└───────────────────────────────────────┴──────────┴──────────┴─────────┴─────────┴────────────┘\n```\n\n## Results PyTorch\n\n```\n┌────────────────────────────────────────────┬──────────┬──────────┬─────────┬─────────┬────────────┐\n│ Name │ Wall/Run │ CPU/Run │ mWd/Run │ Speedup │ vs Fastest │\n├────────────────────────────────────────────┼──────────┼──────────┼─────────┼─────────┼────────────┤\n│ Conv2d B1 C3->32 64x64 K3 f32 (PyTorch) │ 85.48µs │ 130.87µs │ 69.73w │ 1.00x │ 100% │\n│ Conv2d B1 C3->32 64x64 K3 f64 (PyTorch) │ 155.03µs │ 221.13µs │ 118.66w │ 0.55x │ 181% │\n│ Conv2d B8 C32->64 32x32 K3 f32 (PyTorch) │ 756.29µs │ 3.34ms │ 1.75kw │ 0.11x │ 885% │\n│ Conv2d B8 C32->64 32x32 K3 f64 (PyTorch) │ 1.06ms │ 5.96ms │ 3.57kw │ 0.08x │ 1245% │\n│ Conv2d B16 C64->128 16x16 K3 f64 (PyTorch) │ 1.16ms │ 7.32ms │ 4.07kw │ 0.07x │ 1360% │\n│ Conv2d B16 C64->128 16x16 K3 f32 (PyTorch) │ 2.92ms │ 8.72ms │ 4.28kw │ 0.03x │ 3415% │\n└────────────────────────────────────────────┴──────────┴──────────┴─────────┴─────────┴────────────┘\n```\n"
},
{
"path": "packages/nx/bench/conv2d/bench_conv2d_nx.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Benchmarks for Nx.correlate and Nx.extract_patches *)\n\nlet backend_name = \"Nx\"\n\n(* correlate: (leading..., spatial...) with kernel of rank K *)\nlet correlate_configs =\n [\n (* (label, input_shape, kernel_shape) *)\n (\"1D 1k\", [| 100 |], [| 5 |]);\n (\"1D 10k batched\", [| 16; 10000 |], [| 5 |]);\n (\"2D 64x64\", [| 64; 64 |], [| 3; 3 |]);\n (\"2D 256x256\", [| 256; 256 |], [| 3; 3 |]);\n (\"2D batch 8x64x64\", [| 8; 64; 64 |], [| 3; 3 |]);\n (\"2D batch 8x256x256\", [| 8; 256; 256 |], [| 3; 3 |]);\n ]\n\nlet extract_patches_configs =\n [\n (* (label, input_shape, kernel_size, stride) *)\n (\"2D 64x64 k3 s1\", [| 1; 1; 64; 64 |], [| 3; 3 |], [| 1; 1 |]);\n (\"2D 64x64 k3 s2\", [| 1; 1; 64; 64 |], [| 3; 3 |], [| 2; 2 |]);\n (\"2D 256x256 k3 s1\", [| 1; 1; 256; 256 |], [| 3; 3 |], [| 1; 1 |]);\n (\"2D 8x3x64x64 k3 s1\", [| 8; 3; 64; 64 |], [| 3; 3 |], [| 1; 1 |]);\n ]\n\nlet build_benchmarks () =\n let benchmarks = ref [] in\n List.iter\n (fun (label, input_shape, kernel_shape) ->\n let x = Nx.rand Nx.Float32 input_shape in\n let k = Nx.rand Nx.Float32 kernel_shape in\n let name = Printf.sprintf \"correlate %s f32 (%s)\" label backend_name in\n benchmarks :=\n Thumper.bench name (fun () -> Nx.correlate x k) :: !benchmarks)\n correlate_configs;\n List.iter\n (fun (label, input_shape, kernel_size, stride) ->\n let x = Nx.rand Nx.Float32 input_shape in\n let k = Array.length kernel_size in\n let dilation = Array.make k 1 in\n let padding = Array.make k (0, 0) in\n let name =\n Printf.sprintf \"extract_patches %s f32 (%s)\" label backend_name\n in\n benchmarks :=\n Thumper.bench name (fun () ->\n Nx.extract_patches ~kernel_size ~stride ~dilation ~padding x)\n :: !benchmarks)\n extract_patches_configs;\n List.rev !benchmarks\n\nlet () =\n let benchmarks = build_benchmarks () in\n Thumper.run \"nx_conv2d\" benchmarks\n"
},
{
"path": "packages/nx/bench/conv2d/bench_conv2d_pytorch.py",
"content": "from __future__ import annotations\n\nimport sys\nfrom pathlib import Path\nfrom typing import Any, Callable, List, Sequence, Tuple\n\nimport torch\nimport torch.nn.functional as F\n\n_SCRIPTS_DIR = Path(__file__).resolve().parent\nwhile not (_SCRIPTS_DIR / \"dune-project\").exists():\n _SCRIPTS_DIR = _SCRIPTS_DIR.parent\n_SCRIPTS_DIR = _SCRIPTS_DIR / \"scripts\"\nif str(_SCRIPTS_DIR) not in sys.path:\n sys.path.insert(0, str(_SCRIPTS_DIR))\n\nimport ubench # type: ignore\n\n\n# Common CNN layer sizes: (batch, in_channels, out_channels, input_size, kernel_size)\nCONFIGS: Sequence[Tuple[int, int, int, int, int]] = (\n (1, 3, 32, 64, 3), # Small: first conv layer, single image\n (8, 32, 64, 32, 3), # Medium: mid-layer, small batch\n (16, 64, 128, 16, 3), # Large: deep layer, larger batch\n)\n\nDTYPES: Sequence[torch.dtype] = (torch.float32, torch.float64)\nBACKEND_NAME = \"PyTorch\"\n\n\ndef _dtype_label(dtype: torch.dtype) -> str:\n if dtype == torch.float32:\n return \"f32\"\n if dtype == torch.float64:\n return \"f64\"\n return str(dtype)\n\n\ndef _benchmark_name(\n op_name: str,\n batch: int,\n in_ch: int,\n out_ch: int,\n img_size: int,\n kernel_size: int,\n dtype: torch.dtype,\n) -> str:\n return (\n f\"{op_name} B{batch} C{in_ch}->{out_ch} {img_size}x{img_size} \"\n f\"K{kernel_size} {_dtype_label(dtype)} ({BACKEND_NAME})\"\n )\n\n\nclass ConvSpec:\n \"\"\"Conv2d operation specification.\"\"\"\n\n def __init__(\n self,\n name: str,\n batch: int,\n in_channels: int,\n out_channels: int,\n img_size: int,\n kernel_size: int,\n ):\n self.name = name\n self.batch = batch\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.img_size = img_size\n self.kernel_size = kernel_size\n\n\ndef create_conv_specs() -> List[ConvSpec]:\n \"\"\"Create conv2d specs from configs.\"\"\"\n return [\n ConvSpec(\"Conv2d\", batch, in_ch, out_ch, img_size, kernel_size)\n for batch, in_ch, out_ch, img_size, kernel_size in CONFIGS\n ]\n\n\ndef build_benchmarks() -> List[Any]:\n \"\"\"Build all conv2d benchmarks.\"\"\"\n benchmarks: List[Any] = []\n specs = create_conv_specs()\n\n torch.manual_seed(0)\n\n for spec in specs:\n for dtype in DTYPES:\n input_shape = (spec.batch, spec.in_channels, spec.img_size, spec.img_size)\n kernel_shape = (spec.out_channels, spec.in_channels, spec.kernel_size, spec.kernel_size)\n\n input_tensor = torch.rand(input_shape, dtype=dtype)\n kernel_tensor = torch.rand(kernel_shape, dtype=dtype)\n\n bench_name = _benchmark_name(\n spec.name, spec.batch, spec.in_channels, spec.out_channels,\n spec.img_size, spec.kernel_size, dtype\n )\n\n def make_fn(inp: torch.Tensor, kern: torch.Tensor) -> Callable[[], None]:\n return lambda: F.conv2d(inp, kern)\n\n benchmarks.append(ubench.bench(bench_name, make_fn(input_tensor, kernel_tensor)))\n\n return benchmarks\n\n\ndef default_config() -> ubench.Config:\n \"\"\"Create default benchmark configuration.\"\"\"\n return (\n ubench.Config.default()\n .time_limit(1.0)\n .warmup(1)\n .min_measurements(5)\n .min_cpu(0.01)\n .geometric_scale(1.3)\n .gc_stabilization(False)\n .build()\n )\n\n\ndef main() -> None:\n \"\"\"Main entry point.\"\"\"\n benchmarks = build_benchmarks()\n config = default_config()\n ubench.run(benchmarks, config=config, output_format=\"pretty\", verbose=False)\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "packages/nx/bench/conv2d/dune",
"content": "(executable\n (name bench_conv2d_nx)\n (libraries nx thumper))\n\n(rule\n (alias runtest)\n (action\n (progn\n (run %{exe:bench_conv2d_nx.exe} -q)\n (diff? nx_conv2d.thumper nx_conv2d.thumper.corrected))))\n"
},
{
"path": "packages/nx/bench/conv2d/nx_conv2d.thumper",
"content": "# thumper baseline\n# version: 1\n# suite_name: nx_conv2d\n# host: 1480401c3b76ed18\n# cpu: Apple M1 Max\n# ocaml: 5.4.1\n# git: 31747323\n# dirty: true\n# command: ./bench_conv2d_nx.exe -q\n\ncorrelate_1d_10k_batched_f32__nx_\talloc_words\t9.310000e+02\t9.310000e+02\t9.310000e+02\t0.000000e+00\t10\t1\ncorrelate_1d_10k_batched_f32__nx_\tcpu_time\t6.759299e-03\t6.742944e-03\t6.781850e-03\t2.877945e-03\t10\t0\ncorrelate_1d_10k_batched_f32__nx_\twall_time\t6.144004e-03\t6.131814e-03\t6.158235e-03\t2.150099e-03\t10\t0\ncorrelate_1d_1k_f32__nx_\talloc_words\t8.160000e+02\t8.160000e+02\t8.160000e+02\t0.000000e+00\t10\t0\ncorrelate_1d_1k_f32__nx_\tcpu_time\t6.460449e-06\t6.450824e-06\t6.471332e-06\t1.587145e-03\t10\t0\ncorrelate_1d_1k_f32__nx_\twall_time\t6.461020e-06\t6.451447e-06\t6.471827e-06\t1.577114e-03\t10\t1\ncorrelate_2d_256x256_f32__nx_\talloc_words\t8.990000e+02\t8.990000e+02\t8.990000e+02\t0.000000e+00\t10\t0\ncorrelate_2d_256x256_f32__nx_\tcpu_time\t4.401137e-03\t4.389369e-03\t4.411961e-03\t2.566629e-03\t10\t0\ncorrelate_2d_256x256_f32__nx_\twall_time\t3.977256e-03\t3.968869e-03\t3.986821e-03\t2.256904e-03\t10\t0\ncorrelate_2d_64x64_f32__nx_\talloc_words\t8.990000e+02\t8.990000e+02\t8.990000e+02\t0.000000e+00\t10\t2\ncorrelate_2d_64x64_f32__nx_\tcpu_time\t4.350262e-04\t4.309908e-04\t4.382260e-04\t8.315777e-03\t10\t3\ncorrelate_2d_64x64_f32__nx_\twall_time\t3.878239e-04\t3.848593e-04\t3.897182e-04\t6.264335e-03\t10\t2\ncorrelate_2d_batch_8x256x256_f32__nx_\talloc_words\t1.016000e+03\t1.016000e+03\t1.016000e+03\t0.000000e+00\t10\t1\ncorrelate_2d_batch_8x256x256_f32__nx_\tcpu_time\t3.718025e-02\t3.711858e-02\t3.722040e-02\t1.369329e-03\t10\t2\ncorrelate_2d_batch_8x256x256_f32__nx_\twall_time\t3.247770e-02\t3.246563e-02\t3.250109e-02\t5.459608e-04\t10\t1\ncorrelate_2d_batch_8x64x64_f32__nx_\talloc_words\t1.016000e+03\t1.016000e+03\t1.016000e+03\t0.000000e+00\t10\t0\ncorrelate_2d_batch_8x64x64_f32__nx_\tcpu_time\t2.536084e-03\t2.526750e-03\t2.546429e-03\t3.879779e-03\t10\t0\ncorrelate_2d_batch_8x64x64_f32__nx_\twall_time\t2.322492e-03\t2.315861e-03\t2.330876e-03\t3.232401e-03\t10\t0\nextract_patches_2d_256x256_k3_s1_f32__nx_\talloc_words\t1.390000e+02\t1.390000e+02\t1.390000e+02\t0.000000e+00\t10\t0\nextract_patches_2d_256x256_k3_s1_f32__nx_\tcpu_time\t7.138257e-04\t7.129127e-04\t7.148739e-04\t1.373699e-03\t10\t1\nextract_patches_2d_256x256_k3_s1_f32__nx_\twall_time\t7.139240e-04\t7.129017e-04\t7.150828e-04\t1.527537e-03\t10\t0\nextract_patches_2d_64x64_k3_s1_f32__nx_\talloc_words\t1.390000e+02\t1.390000e+02\t1.390000e+02\t0.000000e+00\t10\t1\nextract_patches_2d_64x64_k3_s1_f32__nx_\tcpu_time\t6.823460e-05\t6.808474e-05\t6.831762e-05\t1.706495e-03\t10\t3\nextract_patches_2d_64x64_k3_s1_f32__nx_\twall_time\t6.821857e-05\t6.813969e-05\t6.832750e-05\t1.376542e-03\t10\t2\nextract_patches_2d_64x64_k3_s2_f32__nx_\talloc_words\t1.390000e+02\t1.390000e+02\t1.390000e+02\t0.000000e+00\t10\t0\nextract_patches_2d_64x64_k3_s2_f32__nx_\tcpu_time\t1.160441e-05\t1.159022e-05\t1.162028e-05\t1.295226e-03\t10\t0\nextract_patches_2d_64x64_k3_s2_f32__nx_\twall_time\t1.160520e-05\t1.158941e-05\t1.162151e-05\t1.382895e-03\t10\t0\nextract_patches_2d_8x3x64x64_k3_s1_f32__nx_\talloc_words\t1.390000e+02\t1.390000e+02\t1.390000e+02\t0.000000e+00\t10\t0\nextract_patches_2d_8x3x64x64_k3_s1_f32__nx_\tcpu_time\t1.006386e-03\t1.005213e-03\t1.007881e-03\t1.325399e-03\t10\t0\nextract_patches_2d_8x3x64x64_k3_s1_f32__nx_\twall_time\t1.006492e-03\t1.005348e-03\t1.008017e-03\t1.325981e-03\t10\t0\n"
},
{
"path": "packages/nx/bench/dune",
"content": "(executable\n (name bench_nx)\n (modules bench_nx)\n (libraries nx thumper str))\n\n(rule\n (alias runtest)\n (action\n (progn\n (run %{exe:bench_nx.exe} -q)\n (diff? nx.thumper nx.thumper.corrected))))\n"
},
{
"path": "packages/nx/bench/einsum/README.md",
"content": "# Nx Einsum Benchmarks\n\nComparative benchmarks of Nx einsum operations against NumPy.\n\n## Results Nx\n\n```\n┌──────────────────────────────────┬──────────┬──────────┬─────────┬─────────┬────────────┐\n│ Name │ Wall/Run │ CPU/Run │ mWd/Run │ Speedup │ vs Fastest │\n├──────────────────────────────────┼──────────┼──────────┼─────────┼─────────┼────────────┤\n│ InnerProduct 50x50 f64 (Nx) │ 1.90μs │ 1.92μs │ 1.26kw │ 1.00x │ 100% │\n│ InnerProduct 100x100 f32 (Nx) │ 1.96μs │ 1.99μs │ 1.26kw │ 0.97x │ 103% │\n│ InnerProduct 100x100 f64 (Nx) │ 1.98μs │ 2.00μs │ 1.26kw │ 0.96x │ 104% │\n│ InnerProduct 50x50 f32 (Nx) │ 2.06μs │ 2.03μs │ 1.26kw │ 0.92x │ 108% │\n│ InnerProduct 200x200 f32 (Nx) │ 2.18μs │ 2.13μs │ 1.26kw │ 0.87x │ 114% │\n│ InnerProduct 512x512 f32 (Nx) │ 2.47μs │ 2.44μs │ 1.26kw │ 0.77x │ 130% │\n│ InnerProduct 200x200 f64 (Nx) │ 2.56μs │ 2.45μs │ 1.26kw │ 0.74x │ 135% │\n│ InnerProduct 512x512 f64 (Nx) │ 2.63μs │ 2.63μs │ 1.26kw │ 0.72x │ 138% │\n│ MatMul 50x50 f32 (Nx) │ 3.22μs │ 3.22μs │ 952.00w │ 0.59x │ 169% │\n│ MatMul 50x50 f64 (Nx) │ 4.93μs │ 4.95μs │ 952.00w │ 0.39x │ 259% │\n│ MatMul 100x100 f32 (Nx) │ 6.86μs │ 6.80μs │ 952.00w │ 0.28x │ 360% │\n│ ContractReduce2 50x50 f64 (Nx) │ 7.29μs │ 7.26μs │ 4.01kw │ 0.26x │ 383% │\n│ ContractReduce2 50x50 f32 (Nx) │ 7.30μs │ 7.31μs │ 4.01kw │ 0.26x │ 383% │\n│ ContractReduce1 50x50 f32 (Nx) │ 7.81μs │ 7.80μs │ 4.01kw │ 0.24x │ 410% │\n│ ContractReduce1 50x50 f64 (Nx) │ 8.01μs │ 8.04μs │ 4.01kw │ 0.24x │ 421% │\n│ IndependentSum 50x50 f64 (Nx) │ 9.77μs │ 9.82μs │ 3.81kw │ 0.19x │ 513% │\n│ IndependentSum 50x50 f32 (Nx) │ 10.06μs │ 10.09μs │ 3.81kw │ 0.19x │ 528% │\n│ BatchMatMul 50x50 f32 (Nx) │ 11.52μs │ 11.49μs │ 2.91kw │ 0.17x │ 605% │\n│ ContractReduce2 100x100 f32 (Nx) │ 14.62μs │ 14.62μs │ 4.01kw │ 0.13x │ 768% │\n│ ContractReduce2 100x100 f64 (Nx) │ 14.96μs │ 14.72μs │ 4.01kw │ 0.13x │ 786% │\n│ ContractReduce1 100x100 f64 (Nx) │ 17.04μs │ 17.04μs │ 4.01kw │ 0.11x │ 895% │\n│ ContractReduce1 100x100 f32 (Nx) │ 17.11μs │ 16.96μs │ 4.01kw │ 0.11x │ 899% │\n│ MatMul 100x100 f64 (Nx) │ 33.03μs │ 33.07μs │ 945.00w │ 0.06x │ 1734% │\n│ BatchMatMul 50x50 f64 (Nx) │ 35.48μs │ 35.42μs │ 2.91kw │ 0.05x │ 1864% │\n│ ContractReduce2 200x200 f32 (Nx) │ 50.95μs │ 50.92μs │ 4.01kw │ 0.04x │ 2676% │\n│ ContractReduce2 200x200 f64 (Nx) │ 53.94μs │ 53.27μs │ 4.01kw │ 0.04x │ 2833% │\n│ ContractReduce1 200x200 f32 (Nx) │ 57.76μs │ 57.50μs │ 4.01kw │ 0.03x │ 3033% │\n│ BatchMatMul 100x100 f32 (Nx) │ 57.95μs │ 57.95μs │ 2.91kw │ 0.03x │ 3044% │\n│ ContractReduce1 200x200 f64 (Nx) │ 58.40μs │ 57.82μs │ 4.01kw │ 0.03x │ 3067% │\n│ MatMul 200x200 f32 (Nx) │ 59.92μs │ 59.24μs │ 945.00w │ 0.03x │ 3147% │\n│ BatchMatMul 100x100 f64 (Nx) │ 127.22μs │ 127.21μs │ 2.91kw │ 0.01x │ 6681% │\n│ MatMul 200x200 f64 (Nx) │ 147.15μs │ 145.63μs │ 945.00w │ 0.01x │ 7728% │\n│ IndependentSum 100x100 f64 (Nx) │ 184.71μs │ 471.85μs │ 3.81kw │ 0.01x │ 9701% │\n│ IndependentSum 200x200 f32 (Nx) │ 185.52μs │ 502.80μs │ 3.81kw │ 0.01x │ 9743% │\n│ IndependentSum 200x200 f64 (Nx) │ 186.57μs │ 508.79μs │ 3.81kw │ 0.01x │ 9799% │\n│ IndependentSum 100x100 f32 (Nx) │ 223.17μs │ 460.18μs │ 3.81kw │ 0.01x │ 11721% │\n│ BatchMatMul 200x200 f32 (Nx) │ 235.95μs │ 234.52μs │ 2.91kw │ 0.01x │ 12392% │\n│ IndependentSum 512x512 f64 (Nx) │ 282.82μs │ 1.01ms │ 3.81kw │ 0.01x │ 14853% │\n│ IndependentSum 512x512 f32 (Nx) │ 283.39μs │ 970.86μs │ 3.81kw │ 0.01x │ 14883% │\n│ MatMul 512x512 f32 (Nx) │ 324.12μs │ 424.46μs │ 945.00w │ 0.01x │ 17022% │\n│ ContractReduce2 512x512 f32 (Nx) │ 414.50μs │ 412.71μs │ 4.01kw │ 0.00x │ 21769% │\n│ BatchMatMul 200x200 f64 (Nx) │ 438.50μs │ 437.15μs │ 2.91kw │ 0.00x │ 23029% │\n│ ContractReduce2 512x512 f64 (Nx) │ 446.00μs │ 441.93μs │ 4.01kw │ 0.00x │ 23423% │\n│ ContractReduce1 512x512 f32 (Nx) │ 448.72μs │ 447.88μs │ 4.01kw │ 0.00x │ 23566% │\n│ ContractReduce1 512x512 f64 (Nx) │ 509.76μs │ 507.39μs │ 4.01kw │ 0.00x │ 26772% │\n│ MatMul 512x512 f64 (Nx) │ 766.18μs │ 1.20ms │ 945.00w │ 0.00x │ 40238% │\n│ BatchMatMul 512x512 f32 (Nx) │ 818.47μs │ 1.29ms │ 2.91kw │ 0.00x │ 42985% │\n│ BatchMatMul 512x512 f64 (Nx) │ 2.48ms │ 4.50ms │ 2.91kw │ 0.00x │ 130422% │\n└──────────────────────────────────┴──────────┴──────────┴─────────┴─────────┴────────────┘\n```\n\n## Results NumPy\n\n```\n┌─────────────────────────────────────┬──────────┬─────────┬─────────┬────────────┐\n│ Name │ Time/Run │ mWd/Run │ Speedup │ vs Fastest │\n├─────────────────────────────────────┼──────────┼─────────┼─────────┼────────────┤\n│ InnerProduct 100x100 f32 (NumPy) │ 1.14µs │ 0.12w │ 1.00x │ 100% │\n│ InnerProduct 50x50 f32 (NumPy) │ 1.14µs │ 0.12w │ 1.00x │ 100% │\n│ InnerProduct 50x50 f64 (NumPy) │ 1.14µs │ 0.12w │ 0.99x │ 101% │\n│ InnerProduct 100x100 f64 (NumPy) │ 1.16µs │ 0.13w │ 0.98x │ 102% │\n│ InnerProduct 200x200 f32 (NumPy) │ 1.16µs │ 0.15w │ 0.98x │ 102% │\n│ InnerProduct 200x200 f64 (NumPy) │ 1.23µs │ 0.17w │ 0.93x │ 108% │\n│ ContractReduce2 50x50 f32 (NumPy) │ 10.83µs │ 0.97w │ 0.11x │ 952% │\n│ ContractReduce2 50x50 f64 (NumPy) │ 15.03µs │ 1.29w │ 0.08x │ 1322% │\n│ ContractReduce1 50x50 f32 (NumPy) │ 16.93µs │ 1.64w │ 0.07x │ 1489% │\n│ MatMul 50x50 f32 (NumPy) │ 20.34µs │ 1.74w │ 0.06x │ 1789% │\n│ MatMul 50x50 f64 (NumPy) │ 26.58µs │ 2.77w │ 0.04x │ 2338% │\n│ ContractReduce1 50x50 f64 (NumPy) │ 27.01µs │ 2.81w │ 0.04x │ 2375% │\n│ ContractReduce2 100x100 f32 (NumPy) │ 57.04µs │ 6.33w │ 0.02x │ 5017% │\n│ ContractReduce2 100x100 f64 (NumPy) │ 94.44µs │ 9.32w │ 0.01x │ 8306% │\n│ MatMul 100x100 f32 (NumPy) │ 102.29µs │ 10.32w │ 0.01x │ 8996% │\n│ ContractReduce1 100x100 f32 (NumPy) │ 104.80µs │ 10.60w │ 0.01x │ 9217% │\n│ BatchMatMul 50x50 f32 (NumPy) │ 108.84µs │ 10.09w │ 0.01x │ 9572% │\n│ BatchMatMul 50x50 f64 (NumPy) │ 134.92µs │ 13.38w │ 0.01x │ 11866% │\n│ MatMul 100x100 f64 (NumPy) │ 161.94µs │ 14.91w │ 0.01x │ 14242% │\n│ ContractReduce1 100x100 f64 (NumPy) │ 222.89µs │ 24.17w │ 0.01x │ 19603% │\n│ ContractReduce2 200x200 f32 (NumPy) │ 483.26µs │ 56.67w │ 0.00x │ 42501% │\n│ BatchMatMul 100x100 f32 (NumPy) │ 518.96µs │ 47.34w │ 0.00x │ 45641% │\n│ MatMul 200x200 f32 (NumPy) │ 733.17µs │ 80.95w │ 0.00x │ 64480% │\n│ BatchMatMul 100x100 f64 (NumPy) │ 762.60µs │ 78.79w │ 0.00x │ 67068% │\n│ ContractReduce2 200x200 f64 (NumPy) │ 843.93µs │ 114.13w │ 0.00x │ 74221% │\n│ ContractReduce1 200x200 f32 (NumPy) │ 937.95µs │ 112.44w │ 0.00x │ 82490% │\n│ MatMul 200x200 f64 (NumPy) │ 1.35ms │ 150.59w │ 0.00x │ 118410% │\n│ ContractReduce1 200x200 f64 (NumPy) │ 2.27ms │ 294.93w │ 0.00x │ 200049% │\n│ BatchMatMul 200x200 f32 (NumPy) │ 3.37ms │ 397.35w │ 0.00x │ 296058% │\n│ BatchMatMul 200x200 f64 (NumPy) │ 5.80ms │ 712.76w │ 0.00x │ 510284% │\n└─────────────────────────────────────┴──────────┴─────────┴─────────┴────────────┘\n```\n"
},
{
"path": "packages/nx/bench/einsum/bench_einsum_numpy.py",
"content": "from __future__ import annotations\n\nimport sys\nfrom pathlib import Path\nfrom typing import Any, Callable, List, Sequence\n\nimport numpy as np\n\n_SCRIPTS_DIR = Path(__file__).resolve().parent\nwhile not (_SCRIPTS_DIR / \"dune-project\").exists():\n _SCRIPTS_DIR = _SCRIPTS_DIR.parent\n_SCRIPTS_DIR = _SCRIPTS_DIR / \"scripts\"\nif str(_SCRIPTS_DIR) not in sys.path:\n sys.path.insert(0, str(_SCRIPTS_DIR))\n\nimport ubench # type: ignore\n\n\nSIZES: Sequence[int] = (50, 100, 200)\nDTYPES: Sequence[np.dtype] = (np.float32, np.float64)\nBACKEND_NAME = \"NumPy\"\n_RNG = np.random.default_rng(seed=0)\n\n\ndef _dtype_label(dtype: np.dtype) -> str:\n if dtype == np.float32:\n return \"f32\"\n if dtype == np.float64:\n return \"f64\"\n return str(dtype)\n\n\ndef _benchmark_name(op_name: str, size: int, dtype: np.dtype) -> str:\n return f\"{op_name} {size}x{size} {_dtype_label(dtype)} ({BACKEND_NAME})\"\n\n\nclass EinsumOp:\n \"\"\"Einsum operation specification.\"\"\"\n\n def __init__(\n self,\n name: str,\n subscripts: str,\n setup: Callable[[int, np.dtype], List[np.ndarray]],\n ):\n self.name = name\n self.subscripts = subscripts\n self.setup = setup\n\n\n# Define common einsum operations to benchmark - covering key use cases\nEINSUM_OPS = [\n EinsumOp(\n \"MatMul\",\n \"ij,jk->ik\",\n lambda size, dtype: [\n _RNG.random((size, size), dtype=dtype),\n _RNG.random((size, size), dtype=dtype),\n ],\n ),\n EinsumOp(\n \"BatchMatMul\",\n \"bij,bjk->bik\",\n lambda size, dtype: [\n _RNG.random((4, size, size), dtype=dtype),\n _RNG.random((4, size, size), dtype=dtype),\n ],\n ),\n EinsumOp(\n \"InnerProduct\",\n \"i,i->\",\n lambda size, dtype: [\n _RNG.random(size, dtype=dtype),\n _RNG.random(size, dtype=dtype),\n ],\n ),\n # Critical contraction-reduction patterns (known to be slow in Raven)\n EinsumOp(\n \"ContractReduce1\",\n \"ij,kj->\",\n lambda size, dtype: [\n _RNG.random((size, size), dtype=dtype),\n _RNG.random((size, size), dtype=dtype),\n ],\n ),\n EinsumOp(\n \"ContractReduce2\",\n \"ij,jk->\",\n lambda size, dtype: [\n _RNG.random((size, size), dtype=dtype),\n _RNG.random((size, size), dtype=dtype),\n ],\n ),\n]\n\n\ndef build_benchmarks() -> List[Any]:\n \"\"\"Build all einsum benchmarks.\"\"\"\n benchmarks: List[Any] = []\n\n for size in SIZES:\n for dtype in DTYPES:\n for op in EINSUM_OPS:\n operands = op.setup(size, dtype)\n bench_name = _benchmark_name(op.name, size, dtype)\n\n # Capture operands in closure\n def make_fn(subscripts: str, arrays: List[np.ndarray]) -> Callable[[], None]:\n return lambda: np.einsum(subscripts, *arrays)\n\n benchmarks.append(ubench.bench(bench_name, make_fn(op.subscripts, operands)))\n\n return benchmarks\n\n\ndef default_config() -> ubench.Config:\n \"\"\"Create default benchmark configuration.\"\"\"\n return (\n ubench.Config.default()\n .time_limit(1.0)\n .warmup(1)\n .min_measurements(5)\n .min_cpu(0.01)\n .geometric_scale(1.3)\n .gc_stabilization(False)\n .build()\n )\n\n\ndef main() -> None:\n \"\"\"Main entry point.\"\"\"\n benchmarks = build_benchmarks()\n config = default_config()\n ubench.run(benchmarks, config=config, output_format=\"pretty\", verbose=False)\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "packages/nx/bench/einsum/bench_einsum_nx.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\n(* Configuration *)\nlet sizes = [ 50; 100; 200; 512 ]\nlet backend_name = \"Nx\"\n\nlet benchmark_name op_name size dtype_label =\n Printf.sprintf \"%s %dx%d %s (%s)\" op_name size size dtype_label backend_name\n\ntype einsum_spec = { name : string; subscripts : string }\n\nlet einsum_specs =\n [\n { name = \"MatMul\"; subscripts = \"ij,jk->ik\" };\n { name = \"BatchMatMul\"; subscripts = \"bij,bjk->bik\" };\n { name = \"InnerProduct\"; subscripts = \"i,i->\" };\n (* Critical contraction-reduction patterns (known to be slow) *)\n { name = \"ContractReduce1\"; subscripts = \"ij,kj->\" };\n { name = \"ContractReduce2\"; subscripts = \"ij,jk->\" };\n (* Independent contraction: no shared axes, sum everything *)\n { name = \"IndependentSum\"; subscripts = \"ab,cd->\" };\n ]\n\nlet setup_f32 spec size =\n match spec.name with\n | \"MatMul\" | \"ContractReduce2\" ->\n let shape = [| size; size |] in\n [ Nx.rand Nx.Float32 shape; Nx.rand Nx.Float32 shape ]\n | \"BatchMatMul\" ->\n let shape = [| 4; size; size |] in\n [ Nx.rand Nx.Float32 shape; Nx.rand Nx.Float32 shape ]\n | \"InnerProduct\" ->\n let shape = [| size |] in\n [ Nx.rand Nx.Float32 shape; Nx.rand Nx.Float32 shape ]\n | \"ContractReduce1\" ->\n let shape = [| size; size |] in\n [ Nx.rand Nx.Float32 shape; Nx.rand Nx.Float32 shape ]\n | \"IndependentSum\" ->\n let shape = [| size; size |] in\n [ Nx.rand Nx.Float32 shape; Nx.rand Nx.Float32 shape ]\n | _ -> failwith (\"Unknown einsum operation: \" ^ spec.name)\n\nlet setup_f64 spec size =\n match spec.name with\n | \"MatMul\" | \"ContractReduce2\" ->\n let shape = [| size; size |] in\n [ Nx.rand Nx.Float64 shape; Nx.rand Nx.Float64 shape ]\n | \"BatchMatMul\" ->\n let shape = [| 4; size; size |] in\n [ Nx.rand Nx.Float64 shape; Nx.rand Nx.Float64 shape ]\n | \"InnerProduct\" ->\n let shape = [| size |] in\n [ Nx.rand Nx.Float64 shape; Nx.rand Nx.Float64 shape ]\n | \"ContractReduce1\" ->\n let shape = [| size; size |] in\n [ Nx.rand Nx.Float64 shape; Nx.rand Nx.Float64 shape ]\n | \"IndependentSum\" ->\n let shape = [| size; size |] in\n [ Nx.rand Nx.Float64 shape; Nx.rand Nx.Float64 shape ]\n | _ -> failwith (\"Unknown einsum operation: \" ^ spec.name)\n\nlet build_benchmarks () =\n let f32_benches = ref [] in\n let f64_benches = ref [] in\n\n List.iter\n (fun size ->\n List.iter\n (fun spec ->\n let operands = setup_f32 spec size |> Array.of_list in\n let bench_name = benchmark_name spec.name size \"f32\" in\n let fn () = Nx.einsum spec.subscripts operands in\n f32_benches := Thumper.bench bench_name fn :: !f32_benches)\n einsum_specs)\n sizes;\n\n List.iter\n (fun size ->\n List.iter\n (fun spec ->\n let operands = setup_f64 spec size |> Array.of_list in\n let bench_name = benchmark_name spec.name size \"f64\" in\n let fn () = Nx.einsum spec.subscripts operands in\n f64_benches := Thumper.bench bench_name fn :: !f64_benches)\n einsum_specs)\n sizes;\n\n [\n Thumper.group \"f32\" (List.rev !f32_benches);\n Thumper.group \"f64\" (List.rev !f64_benches);\n ]\n\nlet () =\n let benchmarks = build_benchmarks () in\n Thumper.run \"nx_einsum\" benchmarks\n"
},
{
"path": "packages/nx/bench/einsum/dune",
"content": "(executable\n (name bench_einsum_nx)\n (modules bench_einsum_nx)\n (libraries nx thumper))\n\n(rule\n (alias runtest)\n (action\n (progn\n (run %{exe:bench_einsum_nx.exe} -q)\n (diff? nx_einsum.thumper nx_einsum.thumper.corrected))))\n"
},
{
"path": "packages/nx/bench/einsum/nx_einsum.thumper",
"content": "# thumper baseline\n# version: 1\n# suite_name: nx_einsum\n# host: 1480401c3b76ed18\n# cpu: Apple M1 Max\n# ocaml: 5.4.1\n# git: 31747323\n# dirty: true\n# command: /Users/tmattio/Workspace/raven/_build/default/packages/nx/bench/einsum/bench_einsum_nx.exe --bless --quick\n\nf32/batchmatmul_100x100_f32__nx_\talloc_words\t1.701000e+03\t1.701000e+03\t1.701000e+03\t0.000000e+00\t5\t0\nf32/batchmatmul_100x100_f32__nx_\tcpu_time\t5.518725e-05\t5.460953e-05\t5.569620e-05\t9.845340e-03\t5\t0\nf32/batchmatmul_100x100_f32__nx_\twall_time\t5.536186e-05\t5.486018e-05\t5.585119e-05\t8.950249e-03\t5\t0\nf32/batchmatmul_200x200_f32__nx_\talloc_words\t1.701000e+03\t1.701000e+03\t1.701000e+03\t0.000000e+00\t5\t0\nf32/batchmatmul_200x200_f32__nx_\tcpu_time\t1.977725e-04\t1.925440e-04\t2.030003e-04\t2.643504e-02\t5\t0\nf32/batchmatmul_200x200_f32__nx_\twall_time\t1.980320e-04\t1.925745e-04\t2.032979e-04\t2.707490e-02\t5\t0\nf32/batchmatmul_50x50_f32__nx_\talloc_words\t1.701000e+03\t1.701000e+03\t1.701000e+03\t0.000000e+00\t5\t0\nf32/batchmatmul_50x50_f32__nx_\tcpu_time\t9.819694e-06\t9.608194e-06\t1.004464e-05\t2.222318e-02\t5\t0\nf32/batchmatmul_50x50_f32__nx_\twall_time\t9.822858e-06\t9.623288e-06\t1.006387e-05\t2.242628e-02\t5\t0\nf32/batchmatmul_512x512_f32__nx_\talloc_words\t1.701000e+03\t1.701000e+03\t1.701000e+03\t0.000000e+00\t5\t0\nf32/batchmatmul_512x512_f32__nx_\tcpu_time\t1.297080e-03\t1.259932e-03\t1.340206e-03\t3.094424e-02\t5\t1\nf32/batchmatmul_512x512_f32__nx_\twall_time\t7.969367e-04\t7.790774e-04\t8.238494e-04\t2.809007e-02\t5\t1\nf32/contractreduce1_100x100_f32__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf32/contractreduce1_100x100_f32__nx_\tcpu_time\t1.594601e-05\t1.561383e-05\t1.628084e-05\t2.091450e-02\t5\t0\nf32/contractreduce1_100x100_f32__nx_\twall_time\t1.609735e-05\t1.569702e-05\t1.652829e-05\t2.581989e-02\t5\t1\nf32/contractreduce1_200x200_f32__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf32/contractreduce1_200x200_f32__nx_\tcpu_time\t5.603130e-05\t5.519197e-05\t5.691415e-05\t1.536804e-02\t5\t0\nf32/contractreduce1_200x200_f32__nx_\twall_time\t5.617537e-05\t5.508807e-05\t5.705600e-05\t1.751602e-02\t5\t0\nf32/contractreduce1_50x50_f32__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf32/contractreduce1_50x50_f32__nx_\tcpu_time\t6.199364e-06\t6.090980e-06\t6.304951e-06\t1.725755e-02\t5\t1\nf32/contractreduce1_50x50_f32__nx_\twall_time\t6.211418e-06\t6.096398e-06\t6.304748e-06\t1.677152e-02\t5\t0\nf32/contractreduce1_512x512_f32__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t1\nf32/contractreduce1_512x512_f32__nx_\tcpu_time\t4.509663e-04\t4.489071e-04\t4.521112e-04\t3.552538e-03\t5\t0\nf32/contractreduce1_512x512_f32__nx_\twall_time\t4.516230e-04\t4.494771e-04\t4.528449e-04\t3.728570e-03\t5\t0\nf32/contractreduce2_100x100_f32__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf32/contractreduce2_100x100_f32__nx_\tcpu_time\t1.377649e-05\t1.358060e-05\t1.393764e-05\t1.295804e-02\t5\t0\nf32/contractreduce2_100x100_f32__nx_\twall_time\t1.381472e-05\t1.361141e-05\t1.401211e-05\t1.450244e-02\t5\t0\nf32/contractreduce2_200x200_f32__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf32/contractreduce2_200x200_f32__nx_\tcpu_time\t5.053761e-05\t5.032101e-05\t5.081128e-05\t4.850615e-03\t5\t0\nf32/contractreduce2_200x200_f32__nx_\twall_time\t5.056575e-05\t5.036507e-05\t5.088464e-05\t5.137586e-03\t5\t0\nf32/contractreduce2_50x50_f32__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf32/contractreduce2_50x50_f32__nx_\tcpu_time\t5.939967e-06\t5.918716e-06\t5.964805e-06\t3.879551e-03\t5\t2\nf32/contractreduce2_50x50_f32__nx_\twall_time\t5.953429e-06\t5.924350e-06\t5.978908e-06\t4.582069e-03\t5\t0\nf32/contractreduce2_512x512_f32__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t1\nf32/contractreduce2_512x512_f32__nx_\tcpu_time\t4.102682e-04\t4.066891e-04\t4.159497e-04\t1.128601e-02\t5\t0\nf32/contractreduce2_512x512_f32__nx_\twall_time\t4.107669e-04\t4.070123e-04\t4.173550e-04\t1.258954e-02\t5\t0\nf32/independentsum_100x100_f32__nx_\talloc_words\t1.699000e+03\t1.699000e+03\t1.699000e+03\t0.000000e+00\t9\t0\nf32/independentsum_100x100_f32__nx_\tcpu_time\t4.854485e-04\t4.696863e-04\t5.062099e-04\t3.761847e-02\t9\t0\nf32/independentsum_100x100_f32__nx_\twall_time\t2.070175e-04\t1.960637e-04\t2.159687e-04\t4.807575e-02\t9\t2\nf32/independentsum_200x200_f32__nx_\talloc_words\t1.699000e+03\t1.699000e+03\t1.699000e+03\t0.000000e+00\t5\t0\nf32/independentsum_200x200_f32__nx_\tcpu_time\t4.956030e-04\t4.876978e-04\t5.040143e-04\t1.646122e-02\t5\t0\nf32/independentsum_200x200_f32__nx_\twall_time\t1.944663e-04\t1.892080e-04\t1.992350e-04\t2.578080e-02\t5\t0\nf32/independentsum_50x50_f32__nx_\talloc_words\t1.699000e+03\t1.699000e+03\t1.699000e+03\t0.000000e+00\t5\t0\nf32/independentsum_50x50_f32__nx_\tcpu_time\t8.687534e-06\t8.653824e-06\t8.723281e-06\t3.997540e-03\t5\t0\nf32/independentsum_50x50_f32__nx_\twall_time\t8.703804e-06\t8.674385e-06\t8.736752e-06\t3.582744e-03\t5\t0\nf32/independentsum_512x512_f32__nx_\talloc_words\t1.699000e+03\t1.699000e+03\t1.699000e+03\t0.000000e+00\t5\t0\nf32/independentsum_512x512_f32__nx_\tcpu_time\t9.743405e-04\t9.667292e-04\t9.815193e-04\t7.589837e-03\t5\t0\nf32/independentsum_512x512_f32__nx_\twall_time\t2.947632e-04\t2.914982e-04\t2.971370e-04\t9.564960e-03\t5\t1\nf32/innerproduct_100x100_f32__nx_\talloc_words\t2.900000e+02\t2.900000e+02\t2.900000e+02\t0.000000e+00\t5\t0\nf32/innerproduct_100x100_f32__nx_\tcpu_time\t1.321850e-06\t1.294614e-06\t1.345927e-06\t1.940965e-02\t5\t0\nf32/innerproduct_100x100_f32__nx_\twall_time\t1.349716e-06\t1.298123e-06\t1.406750e-06\t4.024086e-02\t5\t1\nf32/innerproduct_200x200_f32__nx_\talloc_words\t2.900000e+02\t2.900000e+02\t2.900000e+02\t0.000000e+00\t5\t0\nf32/innerproduct_200x200_f32__nx_\tcpu_time\t1.384659e-06\t1.369420e-06\t1.400144e-06\t1.109454e-02\t5\t0\nf32/innerproduct_200x200_f32__nx_\twall_time\t1.386790e-06\t1.375480e-06\t1.399783e-06\t8.762314e-03\t5\t0\nf32/innerproduct_50x50_f32__nx_\talloc_words\t2.900000e+02\t2.900000e+02\t2.900000e+02\t0.000000e+00\t5\t0\nf32/innerproduct_50x50_f32__nx_\tcpu_time\t1.204197e-06\t1.191251e-06\t1.222291e-06\t1.288819e-02\t5\t0\nf32/innerproduct_50x50_f32__nx_\twall_time\t1.204623e-06\t1.190244e-06\t1.222792e-06\t1.350979e-02\t5\t0\nf32/innerproduct_512x512_f32__nx_\talloc_words\t2.900000e+02\t2.900000e+02\t2.900000e+02\t0.000000e+00\t5\t0\nf32/innerproduct_512x512_f32__nx_\tcpu_time\t1.769961e-06\t1.762406e-06\t1.778535e-06\t4.556398e-03\t5\t1\nf32/innerproduct_512x512_f32__nx_\twall_time\t1.771485e-06\t1.764424e-06\t1.779815e-06\t4.344216e-03\t5\t2\nf32/matmul_100x100_f32__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf32/matmul_100x100_f32__nx_\tcpu_time\t6.162875e-06\t6.051695e-06\t6.226873e-06\t1.421240e-02\t5\t2\nf32/matmul_100x100_f32__nx_\twall_time\t6.176118e-06\t6.062691e-06\t6.243488e-06\t1.463678e-02\t5\t2\nf32/matmul_200x200_f32__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf32/matmul_200x200_f32__nx_\tcpu_time\t5.209572e-05\t5.140853e-05\t5.301647e-05\t1.543257e-02\t5\t0\nf32/matmul_200x200_f32__nx_\twall_time\t5.221197e-05\t5.152504e-05\t5.300514e-05\t1.417399e-02\t5\t0\nf32/matmul_50x50_f32__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf32/matmul_50x50_f32__nx_\tcpu_time\t1.785309e-06\t1.751060e-06\t1.823101e-06\t2.017611e-02\t5\t0\nf32/matmul_50x50_f32__nx_\twall_time\t1.785520e-06\t1.751599e-06\t1.826631e-06\t2.101120e-02\t5\t0\nf32/matmul_512x512_f32__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf32/matmul_512x512_f32__nx_\tcpu_time\t3.733167e-04\t3.698187e-04\t3.768665e-04\t9.439403e-03\t5\t0\nf32/matmul_512x512_f32__nx_\twall_time\t2.745620e-04\t2.727650e-04\t2.763596e-04\t6.545934e-03\t5\t0\nf64/batchmatmul_100x100_f64__nx_\talloc_words\t1.701000e+03\t1.701000e+03\t1.701000e+03\t0.000000e+00\t5\t0\nf64/batchmatmul_100x100_f64__nx_\tcpu_time\t1.093644e-04\t1.083188e-04\t1.105216e-04\t1.007089e-02\t5\t0\nf64/batchmatmul_100x100_f64__nx_\twall_time\t1.094756e-04\t1.082373e-04\t1.107356e-04\t1.141050e-02\t5\t0\nf64/batchmatmul_200x200_f64__nx_\talloc_words\t1.701000e+03\t1.701000e+03\t1.701000e+03\t0.000000e+00\t5\t0\nf64/batchmatmul_200x200_f64__nx_\tcpu_time\t3.139812e-04\t3.106339e-04\t3.180848e-04\t1.186526e-02\t5\t0\nf64/batchmatmul_200x200_f64__nx_\twall_time\t3.141270e-04\t3.106870e-04\t3.185770e-04\t1.255869e-02\t5\t0\nf64/batchmatmul_50x50_f64__nx_\talloc_words\t1.701000e+03\t1.701000e+03\t1.701000e+03\t0.000000e+00\t5\t0\nf64/batchmatmul_50x50_f64__nx_\tcpu_time\t3.042245e-05\t3.016933e-05\t3.068369e-05\t8.453563e-03\t5\t0\nf64/batchmatmul_50x50_f64__nx_\twall_time\t3.042087e-05\t3.017943e-05\t3.064479e-05\t7.648615e-03\t5\t0\nf64/batchmatmul_512x512_f64__nx_\talloc_words\t1.701000e+03\t1.701000e+03\t1.701000e+03\t0.000000e+00\t5\t0\nf64/batchmatmul_512x512_f64__nx_\tcpu_time\t4.545035e-03\t4.510594e-03\t4.573573e-03\t6.928353e-03\t5\t0\nf64/batchmatmul_512x512_f64__nx_\twall_time\t2.484462e-03\t2.456361e-03\t2.499965e-03\t8.775377e-03\t5\t1\nf64/contractreduce1_100x100_f64__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf64/contractreduce1_100x100_f64__nx_\tcpu_time\t1.557706e-05\t1.550447e-05\t1.564236e-05\t4.426011e-03\t5\t1\nf64/contractreduce1_100x100_f64__nx_\twall_time\t1.559369e-05\t1.551448e-05\t1.566053e-05\t4.683153e-03\t5\t1\nf64/contractreduce1_200x200_f64__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf64/contractreduce1_200x200_f64__nx_\tcpu_time\t5.395237e-05\t5.369265e-05\t5.411275e-05\t3.893179e-03\t5\t1\nf64/contractreduce1_200x200_f64__nx_\twall_time\t5.396311e-05\t5.369013e-05\t5.412146e-05\t3.996517e-03\t5\t1\nf64/contractreduce1_50x50_f64__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf64/contractreduce1_50x50_f64__nx_\tcpu_time\t6.248360e-06\t6.184944e-06\t6.325221e-06\t1.122510e-02\t5\t0\nf64/contractreduce1_50x50_f64__nx_\twall_time\t6.253809e-06\t6.190696e-06\t6.330196e-06\t1.115326e-02\t5\t0\nf64/contractreduce1_512x512_f64__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t1\nf64/contractreduce1_512x512_f64__nx_\tcpu_time\t4.840624e-04\t4.812912e-04\t4.858375e-04\t4.695958e-03\t5\t0\nf64/contractreduce1_512x512_f64__nx_\twall_time\t4.840321e-04\t4.809812e-04\t4.858114e-04\t4.989489e-03\t5\t0\nf64/contractreduce2_100x100_f64__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf64/contractreduce2_100x100_f64__nx_\tcpu_time\t1.337326e-05\t1.330468e-05\t1.341755e-05\t4.219924e-03\t5\t1\nf64/contractreduce2_100x100_f64__nx_\twall_time\t1.340222e-05\t1.332199e-05\t1.346779e-05\t5.439589e-03\t5\t0\nf64/contractreduce2_200x200_f64__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf64/contractreduce2_200x200_f64__nx_\tcpu_time\t4.878437e-05\t4.843472e-05\t4.914975e-05\t7.328470e-03\t5\t0\nf64/contractreduce2_200x200_f64__nx_\twall_time\t4.879236e-05\t4.848573e-05\t4.915805e-05\t6.889622e-03\t5\t0\nf64/contractreduce2_50x50_f64__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t5\t0\nf64/contractreduce2_50x50_f64__nx_\tcpu_time\t5.661167e-06\t5.639267e-06\t5.688871e-06\t4.381094e-03\t5\t2\nf64/contractreduce2_50x50_f64__nx_\twall_time\t5.669255e-06\t5.642517e-06\t5.701602e-06\t5.210925e-03\t5\t2\nf64/contractreduce2_512x512_f64__nx_\talloc_words\t1.813000e+03\t1.813000e+03\t1.813000e+03\t0.000000e+00\t8\t0\nf64/contractreduce2_512x512_f64__nx_\tcpu_time\t4.410580e-04\t4.296686e-04\t4.590414e-04\t3.329810e-02\t8\t0\nf64/contractreduce2_512x512_f64__nx_\twall_time\t4.411046e-04\t4.302303e-04\t4.575018e-04\t3.091271e-02\t8\t0\nf64/independentsum_100x100_f64__nx_\talloc_words\t1.699000e+03\t1.699000e+03\t1.699000e+03\t0.000000e+00\t5\t0\nf64/independentsum_100x100_f64__nx_\tcpu_time\t4.113704e-04\t4.017748e-04\t4.217268e-04\t2.425059e-02\t5\t0\nf64/independentsum_100x100_f64__nx_\twall_time\t1.732455e-04\t1.688961e-04\t1.779039e-04\t2.599720e-02\t5\t2\nf64/independentsum_200x200_f64__nx_\talloc_words\t1.699000e+03\t1.699000e+03\t1.699000e+03\t0.000000e+00\t5\t0\nf64/independentsum_200x200_f64__nx_\tcpu_time\t4.733798e-04\t4.692991e-04\t4.784865e-04\t9.704117e-03\t5\t0\nf64/independentsum_200x200_f64__nx_\twall_time\t1.991224e-04\t1.954485e-04\t2.025067e-04\t1.772320e-02\t5\t1\nf64/independentsum_50x50_f64__nx_\talloc_words\t1.699000e+03\t1.699000e+03\t1.699000e+03\t0.000000e+00\t5\t0\nf64/independentsum_50x50_f64__nx_\tcpu_time\t8.313444e-06\t8.269508e-06\t8.365449e-06\t5.770216e-03\t5\t2\nf64/independentsum_50x50_f64__nx_\twall_time\t8.318378e-06\t8.276324e-06\t8.372429e-06\t5.776654e-03\t5\t2\nf64/independentsum_512x512_f64__nx_\talloc_words\t1.699000e+03\t1.699000e+03\t1.699000e+03\t0.000000e+00\t5\t1\nf64/independentsum_512x512_f64__nx_\tcpu_time\t9.901247e-04\t9.801404e-04\t1.003114e-03\t1.160142e-02\t5\t0\nf64/independentsum_512x512_f64__nx_\twall_time\t3.018494e-04\t2.935635e-04\t3.142709e-04\t3.430086e-02\t5\t0\nf64/innerproduct_100x100_f64__nx_\talloc_words\t2.900000e+02\t2.900000e+02\t2.900000e+02\t0.000000e+00\t5\t0\nf64/innerproduct_100x100_f64__nx_\tcpu_time\t1.254726e-06\t1.244257e-06\t1.268506e-06\t9.663129e-03\t5\t1\nf64/innerproduct_100x100_f64__nx_\twall_time\t1.254991e-06\t1.245116e-06\t1.267312e-06\t8.843216e-03\t5\t1\nf64/innerproduct_200x200_f64__nx_\talloc_words\t2.900000e+02\t2.900000e+02\t2.900000e+02\t0.000000e+00\t5\t0\nf64/innerproduct_200x200_f64__nx_\tcpu_time\t1.372800e-06\t1.366217e-06\t1.378100e-06\t4.327955e-03\t5\t0\nf64/innerproduct_200x200_f64__nx_\twall_time\t1.373272e-06\t1.367614e-06\t1.379922e-06\t4.481246e-03\t5\t0\nf64/innerproduct_50x50_f64__nx_\talloc_words\t2.900000e+02\t2.900000e+02\t2.900000e+02\t0.000000e+00\t5\t0\nf64/innerproduct_50x50_f64__nx_\tcpu_time\t1.169241e-06\t1.165243e-06\t1.177169e-06\t5.099995e-03\t5\t1\nf64/innerproduct_50x50_f64__nx_\twall_time\t1.169412e-06\t1.165101e-06\t1.177018e-06\t5.095543e-03\t5\t1\nf64/innerproduct_512x512_f64__nx_\talloc_words\t2.900000e+02\t2.900000e+02\t2.900000e+02\t0.000000e+00\t5\t0\nf64/innerproduct_512x512_f64__nx_\tcpu_time\t1.794112e-06\t1.775288e-06\t1.820588e-06\t1.262473e-02\t5\t1\nf64/innerproduct_512x512_f64__nx_\twall_time\t1.796327e-06\t1.776489e-06\t1.823544e-06\t1.309738e-02\t5\t1\nf64/matmul_100x100_f64__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf64/matmul_100x100_f64__nx_\tcpu_time\t2.720731e-05\t2.716254e-05\t2.725049e-05\t1.616393e-03\t5\t0\nf64/matmul_100x100_f64__nx_\twall_time\t2.721241e-05\t2.717116e-05\t2.725219e-05\t1.488889e-03\t5\t0\nf64/matmul_200x200_f64__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf64/matmul_200x200_f64__nx_\tcpu_time\t1.192722e-04\t1.190362e-04\t1.195069e-04\t1.973204e-03\t5\t0\nf64/matmul_200x200_f64__nx_\twall_time\t1.192726e-04\t1.190383e-04\t1.194766e-04\t1.837440e-03\t5\t0\nf64/matmul_50x50_f64__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf64/matmul_50x50_f64__nx_\tcpu_time\t3.063932e-06\t3.014326e-06\t3.101383e-06\t1.420676e-02\t5\t1\nf64/matmul_50x50_f64__nx_\twall_time\t3.065170e-06\t3.015116e-06\t3.105687e-06\t1.477417e-02\t5\t1\nf64/matmul_512x512_f64__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t1\nf64/matmul_512x512_f64__nx_\tcpu_time\t9.722075e-04\t9.688681e-04\t9.791965e-04\t5.311801e-03\t5\t1\nf64/matmul_512x512_f64__nx_\twall_time\t5.703058e-04\t5.670109e-04\t5.757878e-04\t7.694888e-03\t5\t0\n"
},
{
"path": "packages/nx/bench/matmul/README.md",
"content": "# Nx MatMul Benchmarks\n\nFocused benchmarks for dense matrix multiplication comparing Nx and NumPy.\n\nWe benchmark four representative shapes (square, tall-skinny, wide, large) in\nboth `f32` and `f64` to mirror the workloads we rely on for Rune and the Nx\nbackend. Each shape is tested in two modes:\n\n- **alloc**: a fresh output buffer is allocated every call (`Nx.matmul a b` /\n `np.matmul(a, b)`)\n- **reuse**: a pre-allocated output buffer is passed in (`Nx.matmul ~out a b` /\n `np.matmul(a, b, out=out)`)\n\nThe reuse variant isolates pure BLAS compute time from allocation overhead.\n\n## Running the Benchmarks\n\n### Nx (OCaml)\n\n```bash\ndune exec nx/bench/matmul/bench_matmul_nx.exe\n```\n\n### NumPy (Python)\n\n```bash\npython nx/bench/matmul/bench_matmul_numpy.py\n```\n\n## Results Nx (OCaml)\n\n```\n┌─────────────────────────────────────────────────────┬──────────┬──────────┬─────────┬─────────┬────────────┐\n│ Name │ Wall/Run │ CPU/Run │ mWd/Run │ Speedup │ vs Fastest │\n├─────────────────────────────────────────────────────┼──────────┼──────────┼─────────┼─────────┼────────────┤\n│ MatMul SquareSmall 64x64 @ 64x64 f32 reuse (Nx) │ 1.43μs │ 1.48μs │ 408.00w │ 1.00x │ 100% │\n│ MatMul SquareSmall 64x64 @ 64x64 f32 (Nx) │ 2.50μs │ 2.49μs │ 688.00w │ 0.57x │ 175% │\n│ MatMul SquareSmall 64x64 @ 64x64 f64 reuse (Nx) │ 2.82μs │ 2.85μs │ 408.00w │ 0.51x │ 198% │\n│ MatMul SquareSmall 64x64 @ 64x64 f64 (Nx) │ 4.31μs │ 4.31μs │ 688.00w │ 0.33x │ 302% │\n│ MatMul Wide 128x256 @ 256x64 f32 reuse (Nx) │ 5.56μs │ 5.55μs │ 408.00w │ 0.26x │ 390% │\n│ MatMul Wide 128x256 @ 256x64 f32 (Nx) │ 6.46μs │ 6.46μs │ 688.00w │ 0.22x │ 453% │\n│ MatMul TallSkinny 256x64 @ 64x256 f32 reuse (Nx) │ 9.13μs │ 9.12μs │ 408.00w │ 0.16x │ 640% │\n│ MatMul Wide 128x256 @ 256x64 f64 reuse (Nx) │ 16.12μs │ 16.14μs │ 408.00w │ 0.09x │ 1130% │\n│ MatMul Wide 128x256 @ 256x64 f64 (Nx) │ 17.51μs │ 17.50μs │ 688.00w │ 0.08x │ 1227% │\n│ MatMul TallSkinny 256x64 @ 64x256 f64 reuse (Nx) │ 28.96μs │ 28.93μs │ 408.00w │ 0.05x │ 2030% │\n│ MatMul TallSkinny 256x64 @ 64x256 f32 (Nx) │ 66.79μs │ 66.82μs │ 681.00w │ 0.02x │ 4682% │\n│ MatMul TallSkinny 256x64 @ 64x256 f64 (Nx) │ 104.87μs │ 104.87μs │ 681.00w │ 0.01x │ 7350% │\n│ MatMul SquareLarge 512x512 @ 512x512 f32 reuse (Nx) │ 142.69μs │ 240.46μs │ 408.00w │ 0.01x │ 10001% │\n│ MatMul SquareLarge 512x512 @ 512x512 f32 (Nx) │ 244.27μs │ 332.73μs │ 681.00w │ 0.01x │ 17122% │\n│ MatMul SquareLarge 512x512 @ 512x512 f64 reuse (Nx) │ 455.47μs │ 865.78μs │ 408.00w │ 0.00x │ 31925% │\n│ MatMul SquareLarge 512x512 @ 512x512 f64 (Nx) │ 558.42μs │ 965.35μs │ 681.00w │ 0.00x │ 39141% │\n└─────────────────────────────────────────────────────┴──────────┴──────────┴─────────┴─────────┴────────────┘\n```\n\n## Results NumPy (Python)\n\n```\n┌────────────────────────────────────────────────────────┬──────────┬──────────┬─────────┬─────────┬────────────┐\n│ Name │ Wall/Run │ CPU/Run │ mWd/Run │ Speedup │ vs Fastest │\n├────────────────────────────────────────────────────────┼──────────┼──────────┼─────────┼─────────┼────────────┤\n│ MatMul SquareSmall 64x64 @ 64x64 f32 (NumPy) │ 1.62µs │ 1.62µs │ 0.15w │ 1.00x │ 100% │\n│ MatMul SquareSmall 64x64 @ 64x64 f32 reuse (NumPy) │ 1.62µs │ 1.62µs │ 0.15w │ 1.00x │ 100% │\n│ MatMul SquareSmall 64x64 @ 64x64 f64 reuse (NumPy) │ 2.97µs │ 2.97µs │ 0.25w │ 0.55x │ 183% │\n│ MatMul SquareSmall 64x64 @ 64x64 f64 (NumPy) │ 3.03µs │ 3.03µs │ 0.25w │ 0.54x │ 187% │\n│ MatMul Wide 128x256 @ 256x64 f32 reuse (NumPy) │ 5.67µs │ 5.67µs │ 0.58w │ 0.29x │ 349% │\n│ MatMul Wide 128x256 @ 256x64 f32 (NumPy) │ 5.67µs │ 5.67µs │ 0.58w │ 0.29x │ 349% │\n│ MatMul TallSkinny 256x64 @ 64x256 f32 reuse (NumPy) │ 9.17µs │ 9.16µs │ 0.95w │ 0.18x │ 565% │\n│ MatMul TallSkinny 256x64 @ 64x256 f32 (NumPy) │ 9.94µs │ 9.94µs │ 0.94w │ 0.16x │ 613% │\n│ MatMul Wide 128x256 @ 256x64 f64 reuse (NumPy) │ 15.82µs │ 15.82µs │ 1.65w │ 0.10x │ 975% │\n│ MatMul Wide 128x256 @ 256x64 f64 (NumPy) │ 16.82µs │ 16.82µs │ 1.65w │ 0.10x │ 1036% │\n│ MatMul TallSkinny 256x64 @ 64x256 f64 reuse (NumPy) │ 28.73µs │ 28.73µs │ 2.77w │ 0.06x │ 1771% │\n│ MatMul TallSkinny 256x64 @ 64x256 f64 (NumPy) │ 29.63µs │ 29.62µs │ 2.73w │ 0.05x │ 1826% │\n│ MatMul SquareLarge 512x512 @ 512x512 f32 reuse (NumPy) │ 140.12µs │ 238.93µs │ 22.98w │ 0.01x │ 8635% │\n│ MatMul SquareLarge 512x512 @ 512x512 f32 (NumPy) │ 142.90µs │ 241.33µs │ 22.51w │ 0.01x │ 8807% │\n│ MatMul SquareLarge 512x512 @ 512x512 f64 (NumPy) │ 458.38µs │ 872.47µs │ 84.53w │ 0.00x │ 28249% │\n│ MatMul SquareLarge 512x512 @ 512x512 f64 reuse (NumPy) │ 458.52µs │ 870.74µs │ 87.59w │ 0.00x │ 28257% │\n└────────────────────────────────────────────────────────┴──────────┴──────────┴─────────┴─────────┴────────────┘\n```\n\n## Comparison (reuse, f32)\n\nPure BLAS compute time (pre-allocated output, f32):\n\n| Shape | Nx | NumPy | Ratio |\n| -------------------------- | -------- | -------- | --------- |\n| SquareSmall 64x64 | 1.43μs | 1.62μs | **0.88x** |\n| Wide 128x256 @ 256x64 | 5.56μs | 5.67μs | **0.98x** |\n| TallSkinny 256x64 @ 64x256 | 9.13μs | 9.17μs | **1.00x** |\n| SquareLarge 512x512 | 142.69μs | 140.12μs | **1.02x** |\n\nWith a pre-allocated output buffer, Nx is at parity with NumPy.\n\n## Notes on allocation overhead\n\nThe alloc variants show a large gap on some shapes, most dramatically TallSkinny\nf32 where Nx takes 66.79μs (alloc) vs 9.13μs (reuse) — nearly 58μs of pure\nallocation overhead for a 256x256 output buffer (256 KB).\n\nNumPy's alloc path barely suffers (9.94μs vs 9.17μs reuse). This is likely because\nPython's memory allocator (pymalloc) recycles recently freed blocks: in a\nbenchmark loop, each iteration frees and immediately re-allocates the same-sized\noutput, so pymalloc returns the same already-faulted virtual pages. No new page\nfaults occur.\n"
},
{
"path": "packages/nx/bench/matmul/bench_matmul_numpy.py",
"content": "from __future__ import annotations\n\nimport sys\nfrom dataclasses import dataclass\nfrom pathlib import Path\nfrom typing import Any, Callable, List, Sequence, Tuple\n\nimport numpy as np\n\n_SCRIPTS_DIR = Path(__file__).resolve().parent\nwhile not (_SCRIPTS_DIR / \"dune-project\").exists():\n _SCRIPTS_DIR = _SCRIPTS_DIR.parent\n_SCRIPTS_DIR = _SCRIPTS_DIR / \"scripts\"\nif str(_SCRIPTS_DIR) not in sys.path:\n sys.path.insert(0, str(_SCRIPTS_DIR))\n\nimport ubench # type: ignore\n\n\nBACKEND_NAME = \"NumPy\"\nDTYPES: Sequence[np.dtype] = (np.float32, np.float64)\n\n\n@dataclass(frozen=True)\nclass MatmulCase:\n \"\"\"Specification for a matrix multiplication benchmark.\"\"\"\n\n name: str\n m: int\n k: int\n n: int\n seed: int\n\n\nCASES: Sequence[MatmulCase] = (\n MatmulCase(\"SquareSmall\", m=64, k=64, n=64, seed=11),\n MatmulCase(\"TallSkinny\", m=256, k=64, n=256, seed=17),\n MatmulCase(\"Wide\", m=128, k=256, n=64, seed=23),\n MatmulCase(\"SquareLarge\", m=512, k=512, n=512, seed=29),\n)\n\n\ndef _dtype_label(dtype: np.dtype) -> str:\n if dtype == np.float32:\n return \"f32\"\n if dtype == np.float64:\n return \"f64\"\n return str(dtype)\n\n\ndef _benchmark_name(case: MatmulCase, dtype: np.dtype, suffix: str = \"\") -> str:\n return f\"MatMul {case.name} {case.m}x{case.k} @ {case.k}x{case.n} {_dtype_label(dtype)}{suffix} ({BACKEND_NAME})\"\n\n\ndef _make_operands(case: MatmulCase, dtype: np.dtype) -> Tuple[np.ndarray, np.ndarray]:\n lhs_rng = np.random.default_rng(seed=case.seed)\n rhs_rng = np.random.default_rng(seed=case.seed + 1)\n lhs = lhs_rng.random((case.m, case.k), dtype=dtype)\n rhs = rhs_rng.random((case.k, case.n), dtype=dtype)\n return lhs, rhs\n\n\ndef build_benchmarks() -> List[Any]:\n \"\"\"Build benchmarks for NumPy matmul.\"\"\"\n benchmarks: List[Any] = []\n\n for case in CASES:\n for dtype in DTYPES:\n lhs, rhs = _make_operands(case, dtype)\n\n def make_fn(a: np.ndarray, b: np.ndarray) -> Callable[[], None]:\n return lambda: np.matmul(a, b)\n\n benchmarks.append(ubench.bench(_benchmark_name(case, dtype), make_fn(lhs, rhs)))\n\n out = np.empty((case.m, case.n), dtype=dtype)\n\n def make_fn_reuse(a: np.ndarray, b: np.ndarray, o: np.ndarray) -> Callable[[], None]:\n return lambda: np.matmul(a, b, out=o)\n\n benchmarks.append(ubench.bench(_benchmark_name(case, dtype, \" reuse\"), make_fn_reuse(lhs, rhs, out)))\n\n return benchmarks\n\n\ndef default_config() -> ubench.Config:\n \"\"\"Create default benchmark configuration.\"\"\"\n return (\n ubench.Config.default()\n .time_limit(1.0)\n .warmup(1)\n .min_measurements(5)\n .min_cpu(0.01)\n .geometric_scale(1.3)\n .gc_stabilization(False)\n .build()\n )\n\n\ndef main() -> None:\n \"\"\"Main entry point.\"\"\"\n benchmarks = build_benchmarks()\n config = default_config()\n ubench.run(benchmarks, config=config, output_format=\"pretty\", verbose=False)\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "packages/nx/bench/matmul/bench_matmul_nx.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nlet backend_name = \"Nx\"\n\ntype matmul_case = { name : string; m : int; k : int; n : int }\n\nlet cases =\n [\n { name = \"SquareSmall\"; m = 64; k = 64; n = 64 };\n { name = \"TallSkinny\"; m = 256; k = 64; n = 256 };\n { name = \"Wide\"; m = 128; k = 256; n = 64 };\n { name = \"SquareLarge\"; m = 512; k = 512; n = 512 };\n ]\n\nlet benchmark_name case dtype_label suffix =\n Printf.sprintf \"MatMul %s %dx%d @ %dx%d %s%s (%s)\" case.name case.m case.k\n case.k case.n dtype_label suffix backend_name\n\nlet setup_operands (type a b) (dtype : (a, b) Nx.dtype) case =\n let lhs = Nx.rand dtype [| case.m; case.k |] in\n let rhs = Nx.rand dtype [| case.k; case.n |] in\n (lhs, rhs)\n\nlet add_case (type a b) benches case (dtype : (a, b) Nx.dtype) dtype_label =\n let lhs, rhs = setup_operands dtype case in\n let name = benchmark_name case dtype_label \"\" in\n let fn () = Nx.matmul lhs rhs in\n benches := Thumper.bench name fn :: !benches\n\nlet build_benchmarks () =\n let f32_benches = ref [] in\n let f64_benches = ref [] in\n List.iter\n (fun case ->\n add_case f32_benches case Nx.Float32 \"f32\";\n add_case f64_benches case Nx.Float64 \"f64\")\n cases;\n [\n Thumper.group \"f32\" (List.rev !f32_benches);\n Thumper.group \"f64\" (List.rev !f64_benches);\n ]\n\nlet () =\n let benchmarks = build_benchmarks () in\n Thumper.run \"nx_matmul\" benchmarks\n"
},
{
"path": "packages/nx/bench/matmul/dune",
"content": "(executable\n (name bench_matmul_nx)\n (modules bench_matmul_nx)\n (libraries nx thumper))\n\n(rule\n (alias runtest)\n (action\n (progn\n (run %{exe:bench_matmul_nx.exe} -q)\n (diff? nx_matmul.thumper nx_matmul.thumper.corrected))))\n"
},
{
"path": "packages/nx/bench/matmul/nx_matmul.thumper",
"content": "# thumper baseline\n# version: 1\n# suite_name: nx_matmul\n# host: 1480401c3b76ed18\n# cpu: Apple M1 Max\n# ocaml: 5.4.1\n# git: 31747323\n# dirty: true\n# command: /Users/tmattio/Workspace/raven/_build/default/packages/nx/bench/matmul/bench_matmul_nx.exe --bless --quick\n\nf32/matmul_squarelarge_512x512___512x512_f32__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf32/matmul_squarelarge_512x512___512x512_f32__nx_\tcpu_time\t3.098033e-04\t3.078516e-04\t3.124807e-04\t7.470928e-03\t5\t1\nf32/matmul_squarelarge_512x512___512x512_f32__nx_\twall_time\t2.106800e-04\t2.089472e-04\t2.131048e-04\t9.867098e-03\t5\t1\nf32/matmul_squaresmall_64x64___64x64_f32__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf32/matmul_squaresmall_64x64___64x64_f32__nx_\tcpu_time\t1.483007e-06\t1.465233e-06\t1.502532e-06\t1.257557e-02\t5\t2\nf32/matmul_squaresmall_64x64___64x64_f32__nx_\twall_time\t1.483088e-06\t1.465796e-06\t1.505107e-06\t1.325325e-02\t5\t2\nf32/matmul_tallskinny_256x64___64x256_f32__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf32/matmul_tallskinny_256x64___64x256_f32__nx_\tcpu_time\t5.986367e-05\t5.930609e-05\t6.047566e-05\t9.768660e-03\t5\t0\nf32/matmul_tallskinny_256x64___64x256_f32__nx_\twall_time\t5.989386e-05\t5.933197e-05\t6.040919e-05\t8.992800e-03\t5\t0\nf32/matmul_wide_128x256___256x64_f32__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf32/matmul_wide_128x256___256x64_f32__nx_\tcpu_time\t5.925620e-06\t5.888825e-06\t5.966923e-06\t6.589886e-03\t5\t0\nf32/matmul_wide_128x256___256x64_f32__nx_\twall_time\t5.929617e-06\t5.893135e-06\t5.978258e-06\t7.177803e-03\t5\t0\nf64/matmul_squarelarge_512x512___512x512_f64__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t1\nf64/matmul_squarelarge_512x512___512x512_f64__nx_\tcpu_time\t9.528020e-04\t9.434341e-04\t9.591325e-04\t8.238030e-03\t5\t0\nf64/matmul_squarelarge_512x512___512x512_f64__nx_\twall_time\t5.317828e-04\t5.251065e-04\t5.359640e-04\t1.020860e-02\t5\t0\nf64/matmul_squaresmall_64x64___64x64_f64__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf64/matmul_squaresmall_64x64___64x64_f64__nx_\tcpu_time\t3.022022e-06\t3.001909e-06\t3.045365e-06\t7.190003e-03\t5\t1\nf64/matmul_squaresmall_64x64___64x64_f64__nx_\twall_time\t3.024805e-06\t3.004257e-06\t3.049405e-06\t7.462840e-03\t5\t0\nf64/matmul_tallskinny_256x64___64x256_f64__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf64/matmul_tallskinny_256x64___64x256_f64__nx_\tcpu_time\t8.711857e-05\t8.639228e-05\t8.780838e-05\t8.127438e-03\t5\t0\nf64/matmul_tallskinny_256x64___64x256_f64__nx_\twall_time\t8.725737e-05\t8.644738e-05\t8.805477e-05\t9.210608e-03\t5\t0\nf64/matmul_wide_128x256___256x64_f64__nx_\talloc_words\t1.360000e+02\t1.360000e+02\t1.360000e+02\t0.000000e+00\t5\t0\nf64/matmul_wide_128x256___256x64_f64__nx_\tcpu_time\t1.791992e-05\t1.777490e-05\t1.808887e-05\t8.760198e-03\t5\t1\nf64/matmul_wide_128x256___256x64_f64__nx_\twall_time\t1.793295e-05\t1.778910e-05\t1.806828e-05\t7.784129e-03\t5\t1\n"
},
{
"path": "packages/nx/bench/nx.thumper",
"content": "# thumper baseline\n# version: 1\n# suite_name: nx\n# host: 1480401c3b76ed18\n# cpu: Apple M1 Max\n# ocaml: 5.4.1\n# git: 31747323\n# dirty: true\n# command: /Users/tmattio/Workspace/raven/_build/default/packages/nx/bench/bench_nx.exe --bless --quick\n\nf32/add_100x100_f32__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf32/add_100x100_f32__nx_\tcpu_time\t2.289842e-06\t2.255045e-06\t2.324444e-06\t1.515354e-02\t5\t0\nf32/add_100x100_f32__nx_\twall_time\t2.292567e-06\t2.259659e-06\t2.327562e-06\t1.480952e-02\t5\t0\nf32/add_200x200_f32__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t17\t0\nf32/add_200x200_f32__nx_\tcpu_time\t1.221056e-04\t1.180869e-04\t1.254806e-04\t3.027581e-02\t17\t0\nf32/add_200x200_f32__nx_\twall_time\t9.132800e-05\t8.955639e-05\t9.291314e-05\t1.837742e-02\t17\t2\nf32/add_500x500_f32__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf32/add_500x500_f32__nx_\tcpu_time\t1.982840e-04\t1.953932e-04\t2.003410e-04\t1.247646e-02\t5\t0\nf32/add_500x500_f32__nx_\twall_time\t1.331384e-04\t1.319710e-04\t1.343860e-04\t9.069246e-03\t5\t2\nf32/add_50x50_f32__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf32/add_50x50_f32__nx_\tcpu_time\t9.139271e-07\t9.030871e-07\t9.326158e-07\t1.615483e-02\t5\t2\nf32/add_50x50_f32__nx_\twall_time\t9.144538e-07\t9.013141e-07\t9.312193e-07\t1.635138e-02\t5\t2\nf32/mul_100x100_f32__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf32/mul_100x100_f32__nx_\tcpu_time\t2.232628e-06\t2.220281e-06\t2.244351e-06\t5.390599e-03\t5\t0\nf32/mul_100x100_f32__nx_\twall_time\t2.232888e-06\t2.220330e-06\t2.246141e-06\t5.779777e-03\t5\t0\nf32/mul_200x200_f32__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t19\t0\nf32/mul_200x200_f32__nx_\tcpu_time\t1.201819e-04\t1.157264e-04\t1.250037e-04\t3.859657e-02\t19\t1\nf32/mul_200x200_f32__nx_\twall_time\t8.937001e-05\t8.702655e-05\t9.206869e-05\t2.820937e-02\t19\t1\nf32/mul_500x500_f32__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf32/mul_500x500_f32__nx_\tcpu_time\t1.998010e-04\t1.955131e-04\t2.045045e-04\t2.250090e-02\t5\t0\nf32/mul_500x500_f32__nx_\twall_time\t1.381852e-04\t1.347488e-04\t1.447954e-04\t3.635207e-02\t5\t2\nf32/mul_50x50_f32__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf32/mul_50x50_f32__nx_\tcpu_time\t9.028346e-07\t8.931068e-07\t9.136291e-07\t1.136551e-02\t5\t0\nf32/mul_50x50_f32__nx_\twall_time\t9.034284e-07\t8.932583e-07\t9.148317e-07\t1.193975e-02\t5\t0\nf32/sum_100x100_f32__nx_\talloc_words\t1.470000e+02\t1.470000e+02\t1.470000e+02\t0.000000e+00\t5\t0\nf32/sum_100x100_f32__nx_\tcpu_time\t1.651748e-04\t1.621466e-04\t1.690628e-04\t2.093597e-02\t5\t0\nf32/sum_100x100_f32__nx_\twall_time\t9.238947e-05\t9.053141e-05\t9.423027e-05\t2.001771e-02\t5\t0\nf32/sum_200x200_f32__nx_\talloc_words\t1.470000e+02\t1.470000e+02\t1.470000e+02\t0.000000e+00\t5\t0\nf32/sum_200x200_f32__nx_\tcpu_time\t1.989311e-04\t1.979527e-04\t2.003839e-04\t6.110793e-03\t5\t1\nf32/sum_200x200_f32__nx_\twall_time\t1.002071e-04\t9.748815e-05\t1.027331e-04\t2.617052e-02\t5\t0\nf32/sum_500x500_f32__nx_\talloc_words\t1.470000e+02\t1.470000e+02\t1.470000e+02\t0.000000e+00\t5\t0\nf32/sum_500x500_f32__nx_\tcpu_time\t4.345019e-04\t4.263330e-04\t4.394754e-04\t1.512357e-02\t5\t2\nf32/sum_500x500_f32__nx_\twall_time\t1.474488e-04\t1.430595e-04\t1.530280e-04\t3.380340e-02\t5\t2\nf32/sum_50x50_f32__nx_\talloc_words\t1.470000e+02\t1.470000e+02\t1.470000e+02\t0.000000e+00\t5\t0\nf32/sum_50x50_f32__nx_\tcpu_time\t2.932343e-06\t2.916853e-06\t2.948537e-06\t5.402607e-03\t5\t0\nf32/sum_50x50_f32__nx_\twall_time\t2.935280e-06\t2.917064e-06\t2.949515e-06\t5.527796e-03\t5\t0\nf32/transpose_100x100_f32__nx_\talloc_words\t9.600000e+01\t9.600000e+01\t9.600000e+01\t0.000000e+00\t5\t0\nf32/transpose_100x100_f32__nx_\tcpu_time\t1.841045e-07\t1.837203e-07\t1.846028e-07\t2.396922e-03\t5\t1\nf32/transpose_100x100_f32__nx_\twall_time\t1.842488e-07\t1.838311e-07\t1.847817e-07\t2.579518e-03\t5\t0\nf32/transpose_200x200_f32__nx_\talloc_words\t9.600000e+01\t9.600000e+01\t9.600000e+01\t0.000000e+00\t5\t0\nf32/transpose_200x200_f32__nx_\tcpu_time\t1.840894e-07\t1.834876e-07\t1.846789e-07\t3.235640e-03\t5\t1\nf32/transpose_200x200_f32__nx_\twall_time\t1.842397e-07\t1.836915e-07\t1.847628e-07\t2.907419e-03\t5\t1\nf32/transpose_500x500_f32__nx_\talloc_words\t9.600000e+01\t9.600000e+01\t9.600000e+01\t0.000000e+00\t5\t0\nf32/transpose_500x500_f32__nx_\tcpu_time\t1.834628e-07\t1.819333e-07\t1.845332e-07\t7.085538e-03\t5\t1\nf32/transpose_500x500_f32__nx_\twall_time\t1.836824e-07\t1.822987e-07\t1.847774e-07\t6.747174e-03\t5\t1\nf32/transpose_50x50_f32__nx_\talloc_words\t9.600000e+01\t9.600000e+01\t9.600000e+01\t0.000000e+00\t5\t0\nf32/transpose_50x50_f32__nx_\tcpu_time\t1.801717e-07\t1.778792e-07\t1.833509e-07\t1.518458e-02\t5\t0\nf32/transpose_50x50_f32__nx_\twall_time\t1.803913e-07\t1.781703e-07\t1.836419e-07\t1.516598e-02\t5\t0\nf64/add_100x100_f64__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf64/add_100x100_f64__nx_\tcpu_time\t2.126338e-05\t2.101478e-05\t2.156238e-05\t1.287656e-02\t5\t0\nf64/add_100x100_f64__nx_\twall_time\t2.133682e-05\t2.108753e-05\t2.165103e-05\t1.320477e-02\t5\t0\nf64/add_200x200_f64__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf64/add_200x200_f64__nx_\tcpu_time\t1.894001e-04\t1.868058e-04\t1.918444e-04\t1.330146e-02\t5\t2\nf64/add_200x200_f64__nx_\twall_time\t1.351414e-04\t1.344452e-04\t1.359930e-04\t5.726566e-03\t5\t1\nf64/add_500x500_f64__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf64/add_500x500_f64__nx_\tcpu_time\t2.356499e-04\t2.299707e-04\t2.426843e-04\t2.697567e-02\t5\t0\nf64/add_500x500_f64__nx_\twall_time\t1.503470e-04\t1.468076e-04\t1.573330e-04\t3.500354e-02\t5\t1\nf64/add_50x50_f64__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf64/add_50x50_f64__nx_\tcpu_time\t1.121478e-06\t1.103626e-06\t1.134863e-06\t1.392706e-02\t5\t1\nf64/add_50x50_f64__nx_\twall_time\t1.160394e-06\t1.130030e-06\t1.201028e-06\t3.059240e-02\t5\t1\nf64/mul_100x100_f64__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf64/mul_100x100_f64__nx_\tcpu_time\t2.125428e-05\t2.111687e-05\t2.142256e-05\t7.191079e-03\t5\t1\nf64/mul_100x100_f64__nx_\twall_time\t2.127835e-05\t2.112211e-05\t2.146393e-05\t8.032123e-03\t5\t1\nf64/mul_200x200_f64__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t17\t0\nf64/mul_200x200_f64__nx_\tcpu_time\t1.814761e-04\t1.743791e-04\t1.874073e-04\t3.589524e-02\t17\t0\nf64/mul_200x200_f64__nx_\twall_time\t1.352870e-04\t1.329161e-04\t1.380020e-04\t1.879655e-02\t17\t1\nf64/mul_500x500_f64__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf64/mul_500x500_f64__nx_\tcpu_time\t2.469418e-04\t2.437919e-04\t2.492588e-04\t1.106920e-02\t5\t1\nf64/mul_500x500_f64__nx_\twall_time\t1.515527e-04\t1.476347e-04\t1.561116e-04\t2.796696e-02\t5\t1\nf64/mul_50x50_f64__nx_\talloc_words\t1.590000e+02\t1.590000e+02\t1.590000e+02\t0.000000e+00\t5\t0\nf64/mul_50x50_f64__nx_\tcpu_time\t1.084578e-06\t1.080511e-06\t1.090346e-06\t4.534014e-03\t5\t1\nf64/mul_50x50_f64__nx_\twall_time\t1.084654e-06\t1.080573e-06\t1.090667e-06\t4.653020e-03\t5\t1\nf64/sum_100x100_f64__nx_\talloc_words\t1.470000e+02\t1.470000e+02\t1.470000e+02\t0.000000e+00\t6\t0\nf64/sum_100x100_f64__nx_\tcpu_time\t2.008794e-04\t1.965453e-04\t2.072756e-04\t2.670825e-02\t6\t0\nf64/sum_100x100_f64__nx_\twall_time\t9.554162e-05\t9.382525e-05\t9.769565e-05\t2.025505e-02\t6\t0\nf64/sum_200x200_f64__nx_\talloc_words\t1.470000e+02\t1.470000e+02\t1.470000e+02\t0.000000e+00\t5\t0\nf64/sum_200x200_f64__nx_\tcpu_time\t2.272642e-04\t2.234083e-04\t2.335007e-04\t2.220419e-02\t5\t1\nf64/sum_200x200_f64__nx_\twall_time\t1.069815e-04\t1.031181e-04\t1.131110e-04\t4.670401e-02\t5\t1\nf64/sum_500x500_f64__nx_\talloc_words\t1.470000e+02\t1.470000e+02\t1.470000e+02\t0.000000e+00\t5\t0\nf64/sum_500x500_f64__nx_\tcpu_time\t4.457657e-04\t4.412075e-04\t4.507980e-04\t1.075731e-02\t5\t0\nf64/sum_500x500_f64__nx_\twall_time\t1.481280e-04\t1.460633e-04\t1.505301e-04\t1.507728e-02\t5\t0\nf64/sum_50x50_f64__nx_\talloc_words\t1.470000e+02\t1.470000e+02\t1.470000e+02\t0.000000e+00\t5\t0\nf64/sum_50x50_f64__nx_\tcpu_time\t2.815102e-06\t2.810154e-06\t2.819670e-06\t1.690112e-03\t5\t1\nf64/sum_50x50_f64__nx_\twall_time\t2.815492e-06\t2.811970e-06\t2.819541e-06\t1.344629e-03\t5\t1\nf64/transpose_100x100_f64__nx_\talloc_words\t9.600000e+01\t9.600000e+01\t9.600000e+01\t0.000000e+00\t5\t0\nf64/transpose_100x100_f64__nx_\tcpu_time\t1.838005e-07\t1.801340e-07\t1.880564e-07\t2.155166e-02\t5\t1\nf64/transpose_100x100_f64__nx_\twall_time\t1.843967e-07\t1.809335e-07\t1.879831e-07\t1.911527e-02\t5\t1\nf64/transpose_200x200_f64__nx_\talloc_words\t9.600000e+01\t9.600000e+01\t9.600000e+01\t0.000000e+00\t5\t0\nf64/transpose_200x200_f64__nx_\tcpu_time\t1.849292e-07\t1.835602e-07\t1.873169e-07\t1.015702e-02\t5\t1\nf64/transpose_200x200_f64__nx_\twall_time\t1.849859e-07\t1.836472e-07\t1.873381e-07\t9.976236e-03\t5\t1\nf64/transpose_500x500_f64__nx_\talloc_words\t9.600000e+01\t9.600000e+01\t9.600000e+01\t0.000000e+00\t5\t0\nf64/transpose_500x500_f64__nx_\tcpu_time\t1.841557e-07\t1.828908e-07\t1.855841e-07\t7.312706e-03\t5\t0\nf64/transpose_500x500_f64__nx_\twall_time\t1.843796e-07\t1.832046e-07\t1.858367e-07\t7.137778e-03\t5\t1\nf64/transpose_50x50_f64__nx_\talloc_words\t9.600000e+01\t9.600000e+01\t9.600000e+01\t0.000000e+00\t5\t0\nf64/transpose_50x50_f64__nx_\tcpu_time\t1.836186e-07\t1.785118e-07\t1.879890e-07\t2.580672e-02\t5\t0\nf64/transpose_50x50_f64__nx_\twall_time\t1.842476e-07\t1.789137e-07\t1.888056e-07\t2.684399e-02\t5\t0\n"
},
{
"path": "packages/nx/doc/01-getting-started.md",
"content": "# Getting Started\n\nThis guide covers installation, data types, array creation, slicing, broadcasting, and basic operations.\n\n## Installation\n\n\n```bash\nopam install nx\n```\n\nOr build from source:\n\n\n```bash\ngit clone https://github.com/raven-ml/raven\ncd raven && dune build packages/nx\n```\n\nAdd to your `dune` file:\n\n\n```dune\n(executable\n (name main)\n (libraries nx))\n```\n\n## Creating Arrays\n\n```ocaml\nopen Nx\n\nlet () =\n (* From explicit values: provide dtype, shape, and flat data *)\n let a = create Float32 [|2; 3|] [|1.; 2.; 3.; 4.; 5.; 6.|] in\n print_data a;\n\n (* Filled arrays *)\n let z = zeros Float32 [|3; 3|] in\n let o = ones Int32 [|5|] in\n let f = full Float64 [|2; 2|] 3.14 in\n ignore (z, o, f);\n\n (* Ranges and sequences *)\n let r = arange Int32 0 10 1 in (* [0, 1, ..., 9] *)\n let l = linspace Float32 0. 1. 5 in (* 5 points in [0, 1] *)\n ignore (r, l);\n\n (* Random arrays *)\n let x = rand Float32 [|3; 4|] in\n let y = randn Float32 [|3; 4|] in\n ignore (x, y);\n\n (* Special matrices *)\n let i = eye Float32 3 in (* 3×3 identity *)\n print_data i\n```\n\n## Data Types\n\nEvery array has a `dtype` that determines its element type. Common dtypes:\n\n| Dtype | OCaml type | Typical use |\n|-------|-----------|-------------|\n| `Float32` | `float` | Neural networks, images |\n| `Float64` | `float` | Scientific computing |\n| `Int32` | `int32` | Integer data, indices |\n| `Int64` | `int64` | Large integers |\n| `Bool` | `bool` | Masks, conditions |\n| `Complex128` | `Complex.t` | Signal processing |\n\nNx does not automatically cast between types. Convert explicitly with `astype`:\n\n```ocaml\nopen Nx\n\nlet () =\n let x = create Int32 [|3|] [|1l; 2l; 3l|] in\n let y = astype Float32 x in\n print_data y (* [1. 2. 3.] as float32 *)\n```\n\n## Array Properties\n\n```ocaml\nopen Nx\n\nlet () =\n let x = rand Float32 [|2; 3; 4|] in\n Printf.printf \"shape: [|%s|]\\n\"\n (Array.to_list (shape x) |> List.map string_of_int |> String.concat \"; \");\n Printf.printf \"ndim: %d\\n\" (ndim x); (* 3 *)\n Printf.printf \"size: %d\\n\" (size x); (* 24 *)\n Printf.printf \"dtype: %s\\n\" (dtype_to_string (dtype x))\n```\n\n## Element-wise Operations\n\nBinary operations work element-wise and support broadcasting:\n\n```ocaml\nopen Nx\n\nlet () =\n let a = create Float32 [|3|] [|1.; 2.; 3.|] in\n let b = create Float32 [|3|] [|4.; 5.; 6.|] in\n\n let _ = add a b in (* [5. 7. 9.] *)\n let _ = mul a b in (* [4. 10. 18.] *)\n let _ = sub a b in (* [-3. -3. -3.] *)\n let _ = div a b in (* [0.25 0.4 0.5] *)\n\n (* Scalar operations *)\n let _ = add a (scalar Float32 10.) in (* [11. 12. 13.] *)\n\n (* Math functions *)\n let _ = sin a in\n let _ = exp a in\n let _ = sqrt (abs a) in\n ()\n```\n\n## Reductions\n\n```ocaml\nopen Nx\n\nlet () =\n let x = create Float32 [|2; 3|] [|1.; 2.; 3.; 4.; 5.; 6.|] in\n\n (* Reduce all elements *)\n Printf.printf \"sum = %.1f\\n\" (item [] (sum x));\n Printf.printf \"mean = %.1f\\n\" (item [] (mean x));\n\n (* Reduce along an axis *)\n let col_sums = sum ~axes:[0] x in (* sum each column *)\n print_data col_sums; (* [5. 7. 9.] *)\n\n let row_sums = sum ~axes:[1] x in (* sum each row *)\n print_data row_sums (* [6. 15.] *)\n```\n\n## Slicing and Indexing\n\n### Basic indexing\n\n```ocaml\nopen Nx\n\nlet () =\n let x = create Int32 [|3; 3|] [|1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l|] in\n\n (* Get a row *)\n let row = get [1] x in (* [4, 5, 6] *)\n print_data row;\n\n (* Get a scalar *)\n let v = item [1; 2] x in (* 6l *)\n Printf.printf \"x[1,2] = %ld\\n\" v\n```\n\n### Advanced slicing\n\n```ocaml\nopen Nx\n\nlet () =\n let x = create Int32 [|4; 4|]\n [|1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l;\n 9l; 10l; 11l; 12l; 13l; 14l; 15l; 16l|] in\n\n (* Range: rows 0 to 2 (exclusive), all columns *)\n let sub = slice [R (0, 2); A] x in\n print_data sub;\n\n (* Single index on one axis, range on another *)\n let row1_cols = slice [I 1; R (0, 3)] x in\n print_data row1_cols;\n\n (* Gather specific indices *)\n let picked = slice [L [0; 3]; L [1; 2]] x in\n print_data picked\n```\n\nIndex types: `I i` (single index), `R (start, stop)` (half-open range), `Rs (start, stop, step)` (strided range), `L indices` (gather), `A` (all), `N` (new axis).\n\n## Broadcasting\n\nOperations automatically broadcast arrays with compatible shapes. Dimensions are aligned from the right, and each pair must be equal or one must be 1:\n\n```ocaml\nopen Nx\n\nlet () =\n let matrix = ones Float32 [|3; 4|] in\n let row = create Float32 [|1; 4|] [|10.; 20.; 30.; 40.|] in\n let result = add matrix row in (* row added to every row *)\n print_data result\n```\n\n## Matrix Multiplication\n\n```ocaml\nopen Nx\n\nlet () =\n let a = rand Float32 [|3; 4|] in\n let b = rand Float32 [|4; 2|] in\n let c = matmul a b in\n Printf.printf \"(%d×%d) × (%d×%d) = (%d×%d)\\n\"\n (dim 0 a) (dim 1 a) (dim 0 b) (dim 1 b) (dim 0 c) (dim 1 c)\n```\n\n## Next Steps\n\n- [Array Operations](/docs/nx/array-operations/) — reshaping, views, joining, transposing\n- [Linear Algebra](/docs/nx/linear-algebra/) — decompositions, solvers, FFT\n- [NumPy Comparison](/docs/nx/numpy-comparison/) — side-by-side reference if you're coming from Python\n"
},
{
"path": "packages/nx/doc/02-array-operations.md",
"content": "# Array Operations\n\nThis guide covers reshaping, broadcasting, joining, slicing, and the view model that underlies Nx's efficiency.\n\n## Views and Copies\n\nMany Nx operations return **views** — tensors that share the underlying buffer with the original but have different shape, strides, or offset. Views are O(1) and allocate no new data.\n\nView-producing operations: `reshape`, `transpose`, `slice`, `squeeze`, `unsqueeze`, `flip`, `get`, `moveaxis`, `swapaxes`.\n\nCopy-producing operations: `contiguous`, `copy`, `concatenate`, `stack`, `pad`, element-wise operations.\n\nUse `is_c_contiguous` to check whether elements are laid out contiguously in row-major order, and `contiguous` to force a copy when needed:\n\n\n```ocaml\nlet t = Nx.transpose x in\nNx.is_c_contiguous t (* often false *)\nlet t' = Nx.contiguous t (* force a contiguous copy *)\n```\n\n## Reshaping\n\n### reshape\n\nChange the shape without changing the data order. The total number of elements must match. Use `-1` to infer one dimension:\n\n```ocaml\nopen Nx\n\nlet () =\n let x = create Int32 [|6|] [|1l; 2l; 3l; 4l; 5l; 6l|] in\n let a = reshape [|2; 3|] x in\n let b = reshape [|3; -1|] x in (* -1 inferred as 2 *)\n print_data a;\n print_data b\n```\n\n### flatten and unflatten\n\n`flatten` collapses dimensions into one. `unflatten` expands a dimension back:\n\n```ocaml\nopen Nx\n\nlet () =\n let x = zeros Float32 [|2; 3; 4|] in\n ignore (flatten x |> shape); (* [|24|] *)\n ignore (flatten ~start_dim:1 x |> shape); (* [|2; 12|] *)\n\n let y = zeros Float32 [|2; 12|] in\n ignore (unflatten 1 [|3; 4|] y |> shape) (* [|2; 3; 4|] *)\n```\n\n### squeeze and unsqueeze\n\nRemove or add dimensions of size 1:\n\n```ocaml\nopen Nx\n\nlet () =\n let x = ones Float32 [|1; 3; 1; 4|] in\n let a = squeeze x in (* [|3; 4|] *)\n let b = squeeze ~axes:[0] x in (* [|3; 1; 4|] *)\n Printf.printf \"squeeze all: %dx%d\\n\" (dim 0 a) (dim 1 a);\n Printf.printf \"squeeze [0]: %dx%dx%d\\n\" (dim 0 b) (dim 1 b) (dim 2 b);\n\n let y = create Float32 [|3|] [|1.; 2.; 3.|] in\n let c = unsqueeze ~axes:[0; 2] y in (* [|1; 3; 1|] *)\n Printf.printf \"unsqueeze: %dx%dx%d\\n\" (dim 0 c) (dim 1 c) (dim 2 c)\n```\n\n## Broadcasting\n\nBinary operations automatically broadcast operands. Dimensions are aligned from the right, and each pair must be equal or one must be 1:\n\n```ocaml\nopen Nx\n\nlet () =\n (* Add a row vector to every row of a matrix *)\n let matrix = ones Float32 [|3; 4|] in\n let row = create Float32 [|1; 4|] [|10.; 20.; 30.; 40.|] in\n let result = add matrix row in\n print_data result;\n\n (* Add a column vector to every column *)\n let col = create Float32 [|3; 1|] [|100.; 200.; 300.|] in\n let result2 = add matrix col in\n print_data result2\n```\n\nYou can also broadcast explicitly:\n\n\n```ocaml\nlet x = Nx.broadcast_to [|3; 3|] (Nx.create Nx.Float32 [|1; 3|] [|1.; 2.; 3.|])\n(* Repeats the row 3 times without copying data *)\n```\n\n### Broadcasting rules\n\nShapes are compatible when, aligned from the right, every dimension pair is either equal or one of them is 1. The result shape takes the maximum at each position.\n\n```\n[| 3; 4|] + [|1; 4|] → [|3; 4|] ✓\n[|2; 3; 4|] + [| 4|] → [|2; 3; 4|] ✓\n[| 3; 4|] + [|3; 1|] → [|3; 4|] ✓\n[| 3|] + [| 4|] → error ✗\n```\n\n## Transposing and Permuting\n\n### transpose\n\nReverse dimensions (no copy):\n\n```ocaml\nopen Nx\n\nlet () =\n let x = create Int32 [|2; 3|] [|1l; 2l; 3l; 4l; 5l; 6l|] in\n let t = transpose x in\n print_data t\n (* [[1, 4],\n [2, 5],\n [3, 6]] *)\n```\n\nSpecify a permutation for higher-rank tensors:\n\n\n```ocaml\n(* Permute [batch; height; width; channels] to [batch; channels; height; width] *)\nlet nhwc_to_nchw x = Nx.transpose ~axes:[0; 3; 1; 2] x\n```\n\n### moveaxis and swapaxes\n\nMove or swap individual dimensions:\n\n\n```ocaml\nNx.moveaxis 0 2 x (* move axis 0 to position 2 *)\nNx.swapaxes 1 2 x (* swap axes 1 and 2 *)\n```\n\n### flip\n\nReverse elements along axes:\n\n\n```ocaml\nNx.flip ~axes:[1] x (* mirror columns *)\nNx.flip x (* reverse all dimensions *)\n```\n\n## Indexing and Slicing\n\n### get\n\nIndex from the outermost dimension inward. Returns a sub-tensor (view):\n\n```ocaml\nopen Nx\n\nlet () =\n let x = create Int32 [|2; 3|] [|1l; 2l; 3l; 4l; 5l; 6l|] in\n let row = get [1] x in (* second row: [4, 5, 6] *)\n print_data row\n```\n\n### item\n\nExtract a scalar value:\n\n\n```ocaml\nlet v = Nx.item [1; 2] matrix (* element at row 1, column 2 *)\n```\n\n### slice\n\nAdvanced indexing with range and index specifications:\n\n```ocaml\nopen Nx\n\nlet () =\n let x = create Int32 [|3; 3|] [|1l; 2l; 3l; 4l; 5l; 6l; 7l; 8l; 9l|] in\n\n (* R (start, stop): half-open range *)\n let rows_0_1 = slice [R (0, 2); A] x in\n print_data rows_0_1;\n\n (* I i: single index (reduces dimension) *)\n let col_1 = slice [A; I 1] x in\n print_data col_1;\n\n (* L [indices]: gather specific indices *)\n let corners = slice [L [0; 2]; L [0; 2]] x in\n print_data corners\n```\n\nIndex types:\n- `I i` — single index (reduces dimension)\n- `R (start, stop)` — half-open range\n- `Rs (start, stop, step)` — strided range\n- `L indices` — gather listed indices\n- `A` — all elements (default for trailing axes)\n- `N` — insert new axis of size 1\n\n## Joining and Splitting\n\n### concatenate\n\nJoin tensors along an existing axis:\n\n```ocaml\nopen Nx\n\nlet () =\n let a = ones Float32 [|2; 3|] in\n let b = zeros Float32 [|2; 3|] in\n let c = concatenate ~axis:0 [a; b] in (* [|4; 3|] *)\n Printf.printf \"concat axis 0: %dx%d\\n\" (dim 0 c) (dim 1 c);\n let d = concatenate ~axis:1 [a; b] in (* [|2; 6|] *)\n Printf.printf \"concat axis 1: %dx%d\\n\" (dim 0 d) (dim 1 d)\n```\n\nShorthands: `vstack` (axis 0), `hstack` (axis 1), `dstack` (axis 2).\n\n### stack\n\nJoin tensors along a **new** axis:\n\n```ocaml\nopen Nx\n\nlet () =\n let a = create Float32 [|3|] [|1.; 2.; 3.|] in\n let b = create Float32 [|3|] [|4.; 5.; 6.|] in\n let c = stack ~axis:0 [a; b] in (* [|2; 3|] *)\n print_data c\n```\n\n### split\n\nSplit a tensor into equal parts along an axis:\n\n\n```ocaml\nlet parts = Nx.split ~axis:0 2 x (* split into 2 along axis 0 *)\n```\n\n## Tiling and Repeating\n\n### tile\n\nReplicate the tensor according to a repeat pattern:\n\n\n```ocaml\n(* Tile a [2; 3] tensor 2x along rows, 3x along columns → [4; 9] *)\nNx.tile [|2; 3|] x\n```\n\n### repeat\n\nRepeat elements along a single axis:\n\n\n```ocaml\n(* Repeat each element 3 times along axis 0 *)\nNx.repeat ~axis:0 3 x\n```\n\n### pad\n\nPad with a constant value:\n\n\n```ocaml\n(* Pad: 1 before and 2 after along axis 0, 0 and 1 along axis 1 *)\nNx.pad [|(1, 2); (0, 1)|] 0. x\n```\n\n## Next Steps\n\n- [Linear Algebra](/docs/nx/linear-algebra/) — matrix operations, decompositions, FFT\n- [Input/Output](/docs/nx/io/) — reading and writing images, npy, npz files\n- [NumPy Comparison](/docs/nx/numpy-comparison/) — side-by-side reference\n"
},
{
"path": "packages/nx/doc/03-linear-algebra.md",
"content": "# Linear Algebra\n\nNx provides a comprehensive linear algebra suite and FFT operations. This guide covers the most commonly used operations.\n\n## Matrix Multiplication\n\n### matmul\n\nGeneral matrix multiplication supporting batched inputs:\n\n```ocaml\nopen Nx\n\nlet () =\n let a = rand Float32 [|3; 4|] in\n let b = rand Float32 [|4; 2|] in\n let c = matmul a b in (* [|3; 2|] *)\n Printf.printf \"result shape: [|%d; %d|]\\n\" (dim 0 c) (dim 1 c)\n```\n\n`matmul` supports batched matrix multiplication: leading dimensions are broadcast.\n\n\n```ocaml\n(* Batched: [|batch; m; k|] × [|batch; k; n|] → [|batch; m; n|] *)\nlet a = Nx.rand Nx.Float32 [|10; 3; 4|] in\nlet b = Nx.rand Nx.Float32 [|10; 4; 2|] in\nlet c = Nx.matmul a b (* [|10; 3; 2|] *)\n```\n\n### Related products\n\n| Function | Purpose |\n|----------|---------|\n| `dot` | Inner product (flattened inputs) |\n| `vdot` | Complex-conjugate inner product |\n| `inner` | Inner product over last axes |\n| `outer` | Outer product of 1-D tensors |\n| `tensordot` | Contraction over specified axes |\n| `einsum` | Einstein summation notation |\n| `kron` | Kronecker product |\n| `cross` | Cross product of 3-element vectors |\n\n### einsum\n\nEinstein summation provides a compact notation for many tensor operations:\n\n\n```ocaml\n(* Matrix multiplication: ij,jk->ik *)\nlet c = Nx.einsum \"ij,jk->ik\" [|a; b|]\n\n(* Batch matrix multiply: bij,bjk->bik *)\nlet c = Nx.einsum \"bij,bjk->bik\" [|a; b|]\n\n(* Trace: ii-> *)\nlet tr = Nx.einsum \"ii->\" [|m|]\n\n(* Transpose: ij->ji *)\nlet t = Nx.einsum \"ij->ji\" [|m|]\n```\n\n## Decompositions\n\n### Cholesky\n\nFactor a symmetric positive-definite matrix: A = L·Lᵀ\n\n\n```ocaml\nlet l = Nx.cholesky a (* lower triangular by default *)\nlet u = Nx.cholesky ~upper:true a (* upper triangular *)\n```\n\n### QR\n\nFactor A = Q·R where Q is orthogonal and R is upper triangular:\n\n\n```ocaml\nlet q, r = Nx.qr a (* reduced by default *)\nlet q_full, r_full = Nx.qr ~mode:`Complete a\n```\n\n### SVD\n\nSingular value decomposition A = U·Σ·Vᵀ:\n\n\n```ocaml\nlet u, s, vt = Nx.svd a\nlet s_only = Nx.svdvals a (* singular values only, more efficient *)\n```\n\n### Eigendecomposition\n\n\n```ocaml\n(* General: returns complex eigenvalues and eigenvectors *)\nlet eigenvalues, eigenvectors = Nx.eig a\nlet eigenvalues_only = Nx.eigvals a\n\n(* Symmetric/Hermitian: returns real eigenvalues *)\nlet eigenvalues, eigenvectors = Nx.eigh a\nlet eigenvalues_only = Nx.eigvalsh a\n```\n\n## Solving Linear Systems\n\n### solve\n\nSolve A·x = b for x:\n\n\n```ocaml\nlet x = Nx.solve a b\n```\n\n### lstsq\n\nLeast-squares solution (for overdetermined systems):\n\n\n```ocaml\nlet x, residuals, rank, sv = Nx.lstsq a b\n```\n\n### inv and pinv\n\nMatrix inverse and pseudo-inverse:\n\n\n```ocaml\nlet a_inv = Nx.inv a (* requires square, non-singular *)\nlet a_pinv = Nx.pinv a (* works for any shape *)\n```\n\n## Norms and Properties\n\n### norm\n\nCompute various matrix and vector norms:\n\n\n```ocaml\n(* Vector norms *)\nlet l2 = Nx.norm v (* L2 by default *)\nlet l1 = Nx.norm ~ord:(`Float 1.) v (* L1 norm *)\nlet linf = Nx.norm ~ord:`Inf v (* max absolute value *)\n\n(* Matrix norms *)\nlet fro = Nx.norm ~ord:`Fro m (* Frobenius norm *)\n\n(* Along specific axes *)\nlet row_norms = Nx.norm ~axis:[1] m (* per-row L2 norm *)\n```\n\n### Other properties\n\n\n```ocaml\nlet d = Nx.det m (* determinant *)\nlet sd = Nx.slogdet m (* sign and log-determinant *)\nlet tr = Nx.trace m (* sum of diagonal elements *)\nlet r = Nx.matrix_rank m (* numerical rank *)\nlet c = Nx.cond m (* condition number *)\nlet diag = Nx.diagonal m (* extract diagonal *)\n```\n\n## FFT\n\nNx provides the full suite of discrete Fourier transforms.\n\n### Basic FFT\n\n\n```ocaml\n(* 1-D complex FFT and inverse *)\nlet spectrum = Nx.fft x\nlet reconstructed = Nx.ifft spectrum\n\n(* 2-D FFT *)\nlet spectrum_2d = Nx.fft2 image\n\n(* N-D FFT *)\nlet spectrum_nd = Nx.fftn ~axes:[0; 1; 2] volume\n```\n\n### Real FFT\n\nFor real-valued inputs, `rfft` is more efficient — it exploits conjugate symmetry and returns only the positive-frequency half:\n\n\n```ocaml\nlet spectrum = Nx.rfft signal (* n/2+1 complex outputs *)\nlet signal_back = Nx.irfft spectrum (* back to real *)\n\nlet spectrum_2d = Nx.rfft2 image\nlet spectrum_nd = Nx.rfftn ~axes:[0; 1] data\n```\n\n### Frequency axes\n\n\n```ocaml\nlet freqs = Nx.fftfreq n (* frequency bins for fft *)\nlet rfreqs = Nx.rfftfreq n (* frequency bins for rfft *)\nlet shifted = Nx.fftshift spectrum (* shift zero-frequency to center *)\n```\n\n## Next Steps\n\n- [Array Operations](/docs/nx/array-operations/) — reshaping, broadcasting, slicing\n- [Input/Output](/docs/nx/io/) — reading and writing files\n- [NumPy Comparison](/docs/nx/numpy-comparison/) — side-by-side reference\n"
},
{
"path": "packages/nx/doc/04-io.md",
"content": "# Input/Output Operations\n\nThe `Nx_io` module provides functions to load and save Nx tensors in various file formats,\nincluding image formats and NumPy formats.\n\n## Features\n\n- **Image I/O**: Load and save images in PNG, JPEG, BMP, TGA, and GIF formats\n- **NumPy .npy format**: Load and save single arrays in NumPy's native format\n- **NumPy .npz archives**: Load and save multiple named arrays in compressed archives\n- **Runtime dtype detection**: Handle arrays with types determined at runtime\n- **Type conversion utilities**: Convert between different numeric types\n\n## Image Operations\n\n### Loading Images\n\nImages can be loaded as uint8 tensors:\n\n\n```ocaml\nopen Nx_io\n\n(* Load as RGB: shape [|height; width; 3|] *)\nlet img = load_image \"photo.png\"\n\n(* Load as grayscale: shape [|height; width|] *)\nlet gray = load_image ~grayscale:true \"photo.png\"\n```\n\n### Saving Images\n\nSave uint8 tensors as images:\n\n\n```ocaml\n(* Save RGB or grayscale based on shape *)\nsave_image \"output.png\" img\n```\n\nSupported shapes:\n- `[|height; width|]` — Grayscale\n- `[|height; width; 1|]` — Grayscale with explicit channel\n- `[|height; width; 3|]` — RGB\n- `[|height; width; 4|]` — RGBA\n\n## NumPy Format Support\n\n### Single Arrays (.npy)\n\nThe .npy format stores a single array with its dtype and shape information:\n\n\n```ocaml\n(* Load array with runtime-detected type *)\nlet P arr = load_npy \"data.npy\"\n(* arr : ('a, 'b) Nx.t *)\n\n(* Convert to specific type *)\nlet float_arr = load_npy \"data.npy\" |> as_float32\n\n(* Save array *)\nsave_npy \"output.npy\" my_array\n```\n\n### Archives (.npz)\n\nThe .npz format stores multiple named arrays in a compressed archive:\n\n\n```ocaml\n(* Load entire archive *)\nlet archive = load_npz \"bundle.npz\"\n\n(* Access specific array *)\nlet () =\n match Hashtbl.find_opt archive \"weights\" with\n | Some (P arr) ->\n let _weights = as_float32 (P arr) in\n (* use weights *)\n ()\n | None -> failwith \"weights not found\"\n\n(* Load single array directly *)\nlet P data = load_npz_member ~name:\"data\" \"bundle.npz\"\n\n(* Save multiple arrays *)\nlet () =\n save_npz \"model.npz\" [\n (\"inputs\", P input_array);\n (\"labels\", P label_array);\n (\"weights\", P weight_array)\n ]\n```\n\n## Packed Arrays and Type Conversions\n\nSince file formats store type information that's only known at runtime,\nloaded arrays are wrapped in the `packed_nx` type:\n\n```ocaml\ntype packed_nx = P : ('a, 'b) Nx.t -> packed_nx\n```\n\nConvert packed arrays to specific types using the provided functions:\n\n\n```ocaml\nlet packed = load_npy \"data.npy\"\nlet float32_array = as_float32 packed\nlet int32_array = as_int32 packed\nlet uint8_array = as_uint8 packed\n```\n\nAvailable conversions:\n- Floating point: `as_float16`, `as_float32`, `as_float64`\n- Signed integers: `as_int8`, `as_int16`, `as_int32`, `as_int64`\n- Unsigned integers: `as_uint8`, `as_uint16`\n- Complex: `as_complex32`, `as_complex64`\n\n## Examples\n\n### Image Processing Pipeline\n\n```ocaml\nopen Nx\nopen Nx_io\n\nlet process_image input_path output_path =\n (* Load image *)\n let img = load_image input_path in\n\n (* Convert to float for processing *)\n let img_float = Nx.astype float32 img in\n\n (* Normalize to [0, 1] *)\n let normalized = Nx.div_s img_float 255.0 in\n\n (* Apply some processing: reduce brightness and add bias *)\n let processed = Nx.add_s (Nx.mul_s normalized 0.8) 0.1 in\n\n (* Convert back to uint8 *)\n let result = Nx.astype uint8 (Nx.clip ~min:0.0 ~max:255.0 (Nx.mul_s processed 255.0)) in\n\n (* Save result *)\n save_image output_path result\n```\n\n### Model Checkpoint Save/Load\n\n\n```ocaml\nlet save_checkpoint ~path ~epoch ~model =\n let weights = Model.get_weights model in\n let optimizer_state = Model.get_optimizer_state model in\n\n save_npz path [\n (\"epoch\", P (Nx.scalar int32 epoch));\n (\"weights\", P weights);\n (\"optimizer_state\", P optimizer_state);\n ]\n\nlet load_checkpoint path =\n let archive = load_npz path in\n let epoch =\n match Hashtbl.find_opt archive \"epoch\" with\n | Some p -> Nx.item [] (as_int32 p)\n | None -> failwith \"epoch not found\"\n in\n let weights =\n match Hashtbl.find_opt archive \"weights\" with\n | Some p -> as_float32 p\n | None -> failwith \"weights not found\"\n in\n (epoch, weights)\n```\n\n## Error Handling\n\nAll I/O operations may raise `Failure` exceptions:\n- File not found or inaccessible\n- Unsupported file format\n- Invalid data format\n- Incompatible array shapes (for `save_image`)\n- Missing archive members (for `load_npz_member`)\n\n## Performance Considerations\n\n- Image loading/saving uses the stb_image libraries (header-only C libraries)\n- NumPy format I/O is implemented in pure OCaml\n- Large .npz archives are loaded entirely into memory\n- For very large datasets, consider loading arrays individually with `load_npz_member`\n"
},
{
"path": "packages/nx/doc/05-numpy-comparison.md",
"content": "# Nx vs NumPy Comparison\n\nThis document compares the Nx library (OCaml) with NumPy (Python), highlighting similarities, differences, and providing equivalent code examples.\n\n- [Nx vs NumPy Comparison](#nx-vs-numpy-comparison)\n - [1. Overview](#1-overview)\n - [2. Array Creation](#2-array-creation)\n - [Basic Array Creation](#basic-array-creation)\n - [Advanced Array Creation](#advanced-array-creation)\n - [3. Array Operations](#3-array-operations)\n - [Basic Operations](#basic-operations)\n - [Array Manipulation](#array-manipulation)\n - [4. Element Access and Slicing](#4-element-access-and-slicing)\n - [5. Statistical Functions](#5-statistical-functions)\n - [6. Linear Algebra](#6-linear-algebra)\n - [7. Broadcasting](#7-broadcasting)\n - [8. Conditional Operations](#8-conditional-operations)\n - [9. Random Number Generation](#9-random-number-generation)\n - [10. Real-World Example: Linear Regression](#10-real-world-example-linear-regression)\n\n## 1. Overview\n\nNx is a numerical computing library for OCaml. It takes heavy inspiration from NumPy and aims to be as familiar as possible to NumPy users. That said, there are some phyilosophical differences between the two.\n\n- **Pure OCaml Implementation:** Nx is fully native OCaml without C bindings. For that reason, it typically doesn't match NumPy's raw performances. But it is not trying to: while we care about performance, we prioritize the local development experience, where performance is not critical. That said, Nx uses a backend architecture under the hood, so it can easily be extended to use C or CUDA backends. This is what libraries like Rune are doing, implementing custom backends for Nx, making them suitable for production use cases.\n- **Portable Compilation:** In return for a pure OCaml implementation, you get to compile Nx to JavaScript, WebAssembly, or even unikernels. Making it suitable for a wide range of application.\n- **Type Safety First:** Nx leverages OCaml's strong type system and doesn't perform automatic type casting between array types. You can still use the `astype` function for explicit type conversions.\n- **Bigarray Foundation:** Built on OCaml's Bigarray, Nx uses uint8 instead of boolean arrays and doesn't support string arrays.\n\nApart from the above, Nx is designed to be as close to NumPy as possible. The broadcasting rules are the same, and most functions behave similarly. If you notice an undocumented difference, please open an issue; it's probably a bug.\n\n## 2. Array Creation\n\n### Basic Array Creation\n\n**Nx:**\n```ocaml\n(* Creating a zeros array *)\nlet zeros = Nx.zeros Nx.float64 [|3; 3|]\n\n(* Creating a ones array *)\nlet ones = Nx.ones Nx.float64 [|3; 3|]\n\n(* Creating an array with a specific value *)\nlet full = Nx.full Nx.float64 [|3; 3|] 5.0\n\n(* Creating a range *)\nlet range = Nx.arange Nx.int32 0 10 1\n\n(* Creating an identity matrix *)\nlet identity = Nx.identity Nx.float64 3\n```\n\n**NumPy:**\n```python\n# Creating a zeros array\nzeros = np.zeros((3, 3))\n\n# Creating a ones array\nones = np.ones((3, 3))\n\n# Creating an array with a specific value\nfull = np.full((3, 3), 5.0)\n\n# Creating a range\nrange_array = np.arange(0, 10, 1)\n\n# Creating an identity matrix\nidentity = np.identity(3)\n```\n\n### Advanced Array Creation\n\n**Nx:**\n```ocaml\n(* Creating from existing data *)\nlet data = [|1.0; 2.0; 3.0; 4.0|]\nlet arr = Nx.create Nx.float64 [|2; 2|] data\n\n(* Creating using a function *)\nlet init_arr = Nx.init Nx.float64 [|3; 3|] (fun idx ->\n float_of_int (idx.(0) + idx.(1)))\n```\n\n**NumPy:**\n```python\n# Creating from existing data\ndata = [1.0, 2.0, 3.0, 4.0]\narr = np.array(data).reshape(2, 2)\n\n# Creating using a function\ninit_arr = np.fromfunction(lambda i, j: i + j, (3, 3))\n```\n\n## 3. Array Operations\n\n### Basic Operations\n\n**Nx:**\n\n```ocaml\n(* Element-wise addition *)\nlet result = Nx.add arr1 arr2\n\n(* Scalar multiplication *)\nlet scaled = Nx.mul_s arr 2.0\n\n(* Matrix multiplication *)\nlet matmul_result = Nx.matmul arr1 arr2\n```\n\n**NumPy:**\n```python\n# Element-wise addition\nresult = arr1 + arr2\n\n# In-place addition\narr1 += arr2\n\n# Scalar multiplication\nscaled = arr * 2.0\n\n# Matrix multiplication\nmatmul_result = arr1 @ arr2 # or np.matmul(arr1, arr2)\n```\n\n### Array Manipulation\n\n**Nx:**\n\n```ocaml\n(* Reshape array *)\nlet reshaped = Nx.reshape [|1; 6|] arr\n\n(* Transpose array *)\nlet transposed = Nx.transpose arr\n\n(* Flatten array *)\nlet flattened = Nx.flatten arr\n\n(* Concatenate arrays *)\nlet concat = Nx.concatenate ~axis:0 [arr1; arr2]\n```\n\n**NumPy:**\n```python\n# Reshape array\nreshaped = arr.reshape(1, 6)\n\n# Transpose array\ntransposed = arr.transpose()\n\n# Flatten array\nflattened = arr.flatten()\n\n# Concatenate arrays\nconcat = np.concatenate([arr1, arr2], axis=0)\n```\n\n## 4. Element Access and Slicing\n\n**Nx:**\n\n```ocaml\n(* Get a single element *)\nlet element = Nx.item [0; 1] arr\n\n(* Set a single element *)\nlet () = Nx.set_item [0; 1] 5.0 arr\n\n(* Get a slice/subarray *)\nlet row = Nx.get [0] arr\n```\n\n**NumPy:**\n```python\n# Get a single element\nelement = arr[0, 1]\n\n# Set a single element\narr[0, 1] = 5.0\n\n# Get a slice/subarray\nslice = arr[0]\n```\n\n## 5. Statistical Functions\n\n**Nx:**\n\n```ocaml\n(* Sum of all elements *)\nlet total = Nx.sum arr\n\n(* Mean of all elements *)\nlet avg = Nx.mean arr\n\n(* Min and max values *)\nlet min_val = Nx.min arr\nlet max_val = Nx.max arr\n\n(* Sum along an axis *)\nlet axis_sum = Nx.sum ~axes:[|0|] arr\n```\n\n**NumPy:**\n```python\n# Sum of all elements\ntotal = np.sum(arr)\n\n# Mean of all elements\navg = np.mean(arr)\n\n# Min and max values\nmin_val = np.min(arr)\nmax_val = np.max(arr)\n\n# Sum along an axis\naxis_sum = np.sum(arr, axis=0)\n```\n\n## 6. Linear Algebra\n\n**Nx:**\n\n```ocaml\n(* Matrix inverse *)\nlet inv_a = Nx.inv a\n\n(* Solve linear system Ax = b *)\nlet x = Nx.solve a b\n\n(* SVD decomposition *)\nlet u, s, vt = Nx.svd a\n\n(* Eigenvalue decomposition *)\nlet eigenvalues, eigenvectors = Nx.eig a\n```\n\n**NumPy:**\n```python\n# Matrix inverse\ninv_a = np.linalg.inv(a)\n\n# Solve linear system Ax = b\nx = np.linalg.solve(a, b)\n\n# SVD decomposition\nu, s, vt = np.linalg.svd(a)\n\n# Eigenvalue decomposition\neigenvalues, eigenvectors = np.linalg.eig(a)\n```\n\n## 7. Broadcasting\n\n**Nx:**\n\n```ocaml\n(* Broadcast a smaller array to match dimensions *)\nlet broadcasted = Nx.broadcast_to [|3; 3|] smaller_arr\n\n(* Broadcasting happens automatically in operations *)\nlet result = Nx.add matrix vector\n```\n\n**NumPy:**\n```python\n# Broadcast a smaller array to match dimensions\nbroadcasted = np.broadcast_to(smaller_arr, (3, 3))\n\n# Broadcasting happens automatically in operations\nresult = matrix + vector\n```\n\n## 8. Conditional Operations\n\n**Nx:**\n\n```ocaml\n(* Create a boolean mask *)\nlet mask = Nx.greater arr (Nx.scalar Nx.float64 0.5)\n\n(* Apply condition with where *)\nlet result = Nx.where mask arr1 arr2\n```\n\n**NumPy:**\n```python\n# Create a boolean mask\nmask = arr > 0.5\n\n# Apply condition with where\nresult = np.where(mask, arr1, arr2)\n```\n\n## 9. Random Number Generation\n\n**Nx:**\n```ocaml\n(* Generate uniform random numbers *)\nlet random = Nx.rand Nx.float64 [|3; 3|]\n\n(* Generate normal distributed random numbers *)\nlet normal = Nx.randn Nx.float64 [|3; 3|]\n\n(* For reproducibility, wrap in Rng.run *)\nlet reproducible =\n Nx.Rng.run ~seed:42 (fun () ->\n Nx.rand Nx.float64 [|3; 3|])\n```\n\n**NumPy:**\n```python\n# Generate uniform random numbers\nrandom = np.random.rand(3, 3)\n\n# Generate normal distributed random numbers\nnormal = np.random.randn(3, 3)\n```\n\n## 10. Real-World Example: Linear Regression\n\n**Nx:**\n```ocaml\n(* Generate sample data *)\nlet x = Nx.linspace Nx.float64 0.0 10.0 100\nlet y =\n Nx.(\n add (mul_s x 2.0)\n (randn Nx.float64 [|100|]))\n\n(* Reshape x for design matrix *)\nlet x_design = Nx.(concatenate ~axis:1 [ones Nx.float64 [|100; 1|];\n reshape [|100; 1|] x])\n\n(* Compute coefficients using normal equation *)\nlet xtx = Nx.matmul (Nx.transpose x_design) x_design\nlet xty = Nx.matmul (Nx.transpose x_design) (Nx.reshape [|100; 1|] y)\nlet coeffs = Nx.solve xtx xty\n\n(* Make predictions *)\nlet y_pred = Nx.matmul x_design coeffs\n```\n\n**NumPy:**\n```python\n# Generate sample data\nx = np.linspace(0, 10, 100)\ny = 2 * x + np.random.randn(100)\n\n# Reshape x for design matrix\nx_design = np.column_stack((np.ones(100), x))\n\n# Compute coefficients using normal equation\nxtx = x_design.T @ x_design\nxty = x_design.T @ y.reshape(100, 1)\ncoeffs = np.linalg.solve(xtx, xty)\n\n# Make predictions\ny_pred = x_design @ coeffs\n```\n"
},
{
"path": "packages/nx/doc/dune",
"content": "(mdx\n (files *.md)\n (package nx)\n (libraries nx nx.io))\n"
},
{
"path": "packages/nx/doc/index.md",
"content": "# nx\n\nNx provides n-dimensional arrays with NumPy-like semantics and OCaml's type safety. It is the numerical foundation for the entire Raven ecosystem.\n\n## Features\n\n- **19 data types** — float16 through float64, int4 through int64, complex64/128, bool\n- **Broadcasting** — automatic shape matching for binary operations\n- **Views** — reshape, transpose, and slice without copying data\n- **Linear algebra** — matmul, solve, cholesky, QR, SVD, eigendecomposition\n- **FFT** — full suite of discrete Fourier transforms\n- **Signal processing** — convolution, correlation, filtering\n- **I/O** — read and write images (PNG, JPEG), NumPy files (.npy, .npz)\n- **Pluggable backends** — default C backend, extensible architecture\n\n## Quick Start\n\n```ocaml\nopen Nx\n\nlet () =\n (* Create and manipulate arrays *)\n let x = linspace Float32 0. 10. 5 in\n let y = mul x x in\n Printf.printf \"x = \"; print_data x;\n Printf.printf \"y = x² = \"; print_data y;\n\n (* Matrix operations *)\n let a = rand Float32 [|3; 3|] in\n let b = rand Float32 [|3; 3|] in\n let c = matmul a b in\n Printf.printf \"matmul shape: [|%d; %d|]\\n\" (dim 0 c) (dim 1 c)\n```\n\n## Next Steps\n\n- [Getting Started](/docs/nx/getting-started/) — installation, dtypes, slicing, broadcasting\n- [Array Operations](/docs/nx/array-operations/) — reshaping, views, joining, splitting\n- [Linear Algebra](/docs/nx/linear-algebra/) — decompositions, solvers, FFT\n- [Input/Output](/docs/nx/io/) — images, npy, npz files\n- [NumPy Comparison](/docs/nx/numpy-comparison/) — side-by-side reference\n"
},
{
"path": "packages/nx/examples/01-creating-arrays/README.md",
"content": "# `01-creating-arrays`\n\nBuild arrays from scratch — constants, ranges, grids, and custom data. This\nexample walks through the most common ways to create arrays in Nx.\n\n```bash\ndune exec nx/examples/01-creating-arrays/main.exe\n```\n\n## What You'll Learn\n\n- Choosing a dtype (`float32`, `float64`, `int32`)\n- Filling arrays with constants: `zeros`, `ones`, `full`\n- Generating ranges: `arange`, `linspace`, `logspace`\n- Building arrays from OCaml data: `create`, `init`\n- Diagonal and special matrices: `identity`, `eye`, `tril`, `triu`\n- Coordinate grids with `meshgrid`\n\n## Key Functions\n\n| Function | Purpose |\n| ------------------------------ | -------------------------------------- |\n| `zeros dtype shape` | Array of all zeros |\n| `ones dtype shape` | Array of all ones |\n| `full dtype shape value` | Array filled with a value |\n| `arange dtype start stop step` | Integer-stepped range (exclusive stop) |\n| `linspace dtype start stop n` | Evenly spaced floats |\n| `logspace dtype start stop n` | Logarithmically spaced values |\n| `create dtype shape data` | Array from an OCaml array |\n| `init dtype shape f` | Array from a function of indices |\n| `identity dtype n` | n×n identity matrix |\n| `eye ?k dtype n` | Ones on the k-th diagonal |\n| `meshgrid x y` | Coordinate grids from 1D arrays |\n| `tril m` / `triu m` | Lower / upper triangular part |\n\n## Output Walkthrough\n\nWhen you run this example, you'll see arrays printed in a compact format:\n\n```\nzeros (2×3):\n[[0, 0, 0],\n [0, 0, 0]]\n\narange 0..9:\n[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]\n\n5×5 multiplication table:\n[[1, 2, 3, 4, 5],\n [2, 4, 6, 8, 10],\n [3, 6, 9, 12, 15],\n [4, 8, 12, 16, 20],\n [5, 10, 15, 20, 25]]\n```\n\nThe multiplication table is built with `init`, which calls a function with the\nindex array for each element:\n\n```ocaml\ninit int32 [| 5; 5 |] (fun idx ->\n Int32.of_int ((idx.(0) + 1) * (idx.(1) + 1)))\n```\n\n`meshgrid` builds a pair of 2D coordinate grids from two 1D arrays — useful\nfor evaluating functions over a grid:\n\n```\nmeshgrid X:\n[[0, 1, 2],\n [0, 1, 2]]\nmeshgrid Y:\n[[0, 0, 0],\n [1, 1, 1]]\n```\n\n## Dtypes\n\nEvery array has a dtype that determines the element type and precision. The\nfirst argument to most creation functions is the dtype:\n\n```ocaml\nzeros float32 [| 2; 3 |] (* 32-bit floats *)\nones float64 [| 3 |] (* 64-bit floats *)\narange int32 0 10 1 (* 32-bit integers *)\n```\n\nNx supports 18 dtypes including `Float16`, `BFloat16`, `Complex128`, `Bool`,\nand various integer widths.\n\n## Try It\n\n1. Create a 10-element `linspace` from -1.0 to 1.0 and print it.\n2. Use `init` to build a 4×4 matrix where each element is the sum of its row\n and column index.\n3. Try `eye ~k:(-1) float64 4` to see the subdiagonal.\n\n## Next Steps\n\nContinue to [02-infix-and-arithmetic](../02-infix-and-arithmetic/) to learn how\nthe Infix module makes array math read like algebra.\n"
},
{
"path": "packages/nx/examples/01-creating-arrays/dune",
"content": "(executable\n (name main)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/01-creating-arrays/main.ml",
"content": "(** Build arrays from scratch — constants, ranges, grids, and custom data.\n\n This example walks through the most common ways to create arrays. By the end\n you'll know how to pick a dtype, fill arrays with constants, generate\n ranges, and build grids and triangular matrices. *)\n\nopen Nx\n\nlet () =\n (* Constant-filled arrays: zeros, ones, and an arbitrary fill value. *)\n let z = zeros float32 [| 2; 3 |] in\n Printf.printf \"zeros (2×3):\\n%s\\n\\n\" (data_to_string z);\n\n let o = ones float64 [| 3 |] in\n Printf.printf \"ones (3):\\n%s\\n\\n\" (data_to_string o);\n\n let pi = full float64 [| 2; 2 |] Float.pi in\n Printf.printf \"full π (2×2):\\n%s\\n\\n\" (data_to_string pi);\n\n (* Ranges: integer steps and evenly-spaced floats. *)\n let ints = arange int32 0 10 1 in\n Printf.printf \"arange 0..9:\\n%s\\n\\n\" (data_to_string ints);\n\n let spaced = linspace float64 0.0 1.0 5 in\n Printf.printf \"linspace 0..1 (5 points):\\n%s\\n\\n\" (data_to_string spaced);\n\n let decades = logspace float64 1.0 4.0 4 in\n Printf.printf \"logspace 10¹..10⁴:\\n%s\\n\\n\" (data_to_string decades);\n\n (* From raw data: pack an OCaml array into a 2×3 matrix. *)\n let data = create float64 [| 2; 3 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0 |] in\n Printf.printf \"create from data (2×3):\\n%s\\n\\n\" (data_to_string data);\n\n (* Build a multiplication table with [init]. *)\n let mul_table =\n init int32 [| 5; 5 |] (fun idx ->\n Int32.of_int ((idx.(0) + 1) * (idx.(1) + 1)))\n in\n Printf.printf \"5×5 multiplication table:\\n%s\\n\\n\" (data_to_string mul_table);\n\n (* Identity and eye: diagonal matrices. *)\n let id = identity float64 3 in\n Printf.printf \"identity 3×3:\\n%s\\n\\n\" (data_to_string id);\n\n let e = eye ~k:1 float64 3 in\n Printf.printf \"eye (k=1, superdiagonal):\\n%s\\n\\n\" (data_to_string e);\n\n (* Coordinate grids with meshgrid. *)\n let xs = arange_f float64 0.0 3.0 1.0 in\n let ys = arange_f float64 0.0 2.0 1.0 in\n let grid_x, grid_y = meshgrid xs ys in\n Printf.printf \"meshgrid X:\\n%s\\n\" (data_to_string grid_x);\n Printf.printf \"meshgrid Y:\\n%s\\n\\n\" (data_to_string grid_y);\n\n (* Triangular matrices: tril and triu. *)\n let m = ones float64 [| 4; 4 |] in\n Printf.printf \"tril (lower triangle):\\n%s\\n\" (data_to_string (tril m));\n Printf.printf \"triu (upper triangle):\\n%s\\n\" (data_to_string (triu m))\n"
},
{
"path": "packages/nx/examples/02-infix-and-arithmetic/README.md",
"content": "# `02-infix-and-arithmetic`\n\nElement-wise math with operators — the Infix module makes array code read like\nalgebra.\n\n```bash\ndune exec nx/examples/02-infix-and-arithmetic/main.exe\n```\n\n## What You'll Learn\n\n- Using the `Infix` module for clean operator-based math\n- Scalar arithmetic: `*$`, `+$`, `-$`, `/$` (array op scalar)\n- Element-wise operations: `*`, `/`, `+`, `-` (array op array)\n- Math functions: `abs`, `sqrt`, `square`, `exp`, `sign`\n- Clamping values with `clamp ~min ~max`\n- Min-max normalization for scaling data to [0, 1]\n\n## Key Functions\n\n| Function / Operator | Purpose |\n| ------------------- | ---------------------------------- |\n| `a +$ s` | Add scalar to all elements |\n| `a -$ s` | Subtract scalar from all elements |\n| `a *$ s` | Multiply all elements by scalar |\n| `a /$ s` | Divide all elements by scalar |\n| `a + b` | Element-wise addition |\n| `a - b` | Element-wise subtraction |\n| `a * b` | Element-wise multiplication |\n| `a / b` | Element-wise division |\n| `abs a` | Absolute value of each element |\n| `sqrt a` | Square root of each element |\n| `square a` | Square each element |\n| `exp a` | e^x for each element |\n| `sign a` | Sign of each element (-1, 0, or 1) |\n| `clamp ~min ~max a` | Clip values to [min, max] |\n| `min a` / `max a` | Minimum / maximum element |\n\n## Output Walkthrough\n\nWhen you run this example, you'll see various arithmetic operations applied to\narrays:\n\n```\nCelsius: [0, 20, 37, 100, -40]\nFahrenheit: [32, 68, 98.6, 212, -40]\n```\n\nTemperature conversion uses scalar arithmetic — `*$` for multiplication and\n`+$` for addition:\n\n```ocaml\nlet fahrenheit = celsius *$ 1.8 +$ 32.0 in\n```\n\nThis reads just like the formula `F = C × 1.8 + 32`, but operates on the entire\narray at once.\n\nBMI calculation demonstrates element-wise array operations:\n\n```\nHeights (m): [1.65, 1.8, 1.72, 1.55]\nWeights (kg): [68, 90, 75, 52]\nBMI: [24.977, 27.778, 25.351, 21.639]\n```\n\nThe formula `BMI = weight / height²` becomes:\n\n```ocaml\nlet bmi = weight_kg / (height_m * height_m) in\n```\n\nBoth `*` and `/` work element-by-element, computing BMI for all individuals\nsimultaneously.\n\nMin-max normalization scales exam scores to the range [0, 1]:\n\n```\nRaw scores: [72, 85, 60, 93, 78, 55]\nNormalized: [0.447, 0.789, 0.132, 1, 0.605, 0]\n```\n\nThe implementation extracts minimum and maximum values, then applies the\nnormalization formula:\n\n```ocaml\nlet lo = min scores in\nlet hi = max scores in\nlet normalized = (scores - lo) / (hi - lo) in\n```\n\nMath functions apply element-wise to transform arrays:\n\n```\nx: [-2, -1, 0, 1, 2]\nabs(x): [2, 1, 0, 1, 2]\nx²: [4, 1, 0, 1, 4]\n√|x|: [1.414, 1, 0, 1, 1.414]\nexp(x): [0.135, 0.368, 1, 2.718, 7.389]\nsign(x): [-1, -1, 0, 1, 1]\n```\n\n`clamp` restricts values to a valid range — useful for sensor data or ensuring\ninputs stay within bounds:\n\n```\nSensor readings: [-5, 12, 105, 42, -1, 99]\nClamped [0,100]: [0, 12, 100, 42, 0, 99]\n```\n\n## Try It\n\n1. Create an array of angles in degrees and convert to radians using `*$`.\n2. Build two 1D arrays and compute their element-wise difference, then use\n `abs` to get absolute differences.\n3. Generate random-looking data with `create`, then normalize it to [-1, 1]\n instead of [0, 1].\n\n## Next Steps\n\nContinue to [03-indexing-and-slicing](../03-indexing-and-slicing/) to learn how\nto extract and modify specific regions of arrays.\n"
},
{
"path": "packages/nx/examples/02-infix-and-arithmetic/dune",
"content": "(executable\n (name main)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/02-infix-and-arithmetic/main.ml",
"content": "(** Element-wise math with operators — the Infix module makes array code read\n like algebra.\n\n Temperature conversions, BMI calculations, and score normalization, all\n expressed with clean infix operators instead of verbose function calls. *)\n\nopen Nx\nopen Nx.Infix\n\nlet () =\n (* --- Temperature conversion: C → F --- *)\n let celsius = create float64 [| 5 |] [| 0.0; 20.0; 37.0; 100.0; -40.0 |] in\n let fahrenheit = (celsius *$ 1.8) +$ 32.0 in\n Printf.printf \"Celsius: %s\\n\" (data_to_string celsius);\n Printf.printf \"Fahrenheit: %s\\n\\n\" (data_to_string fahrenheit);\n\n (* --- BMI from height and weight arrays --- *)\n let height_m = create float64 [| 4 |] [| 1.65; 1.80; 1.72; 1.55 |] in\n let weight_kg = create float64 [| 4 |] [| 68.0; 90.0; 75.0; 52.0 |] in\n let bmi = weight_kg / (height_m * height_m) in\n Printf.printf \"Heights (m): %s\\n\" (data_to_string height_m);\n Printf.printf \"Weights (kg): %s\\n\" (data_to_string weight_kg);\n Printf.printf \"BMI: %s\\n\\n\" (data_to_string bmi);\n\n (* --- Exam score normalization (min-max scaling to [0, 1]) --- *)\n let scores =\n create float64 [| 6 |] [| 72.0; 85.0; 60.0; 93.0; 78.0; 55.0 |]\n in\n let lo = min scores in\n let hi = max scores in\n let normalized = (scores - lo) / (hi - lo) in\n Printf.printf \"Raw scores: %s\\n\" (data_to_string scores);\n Printf.printf \"Normalized: %s\\n\\n\" (data_to_string normalized);\n\n (* --- Math functions: exp, log, sqrt, abs --- *)\n let x = create float64 [| 5 |] [| -2.0; -1.0; 0.0; 1.0; 2.0 |] in\n Printf.printf \"x: %s\\n\" (data_to_string x);\n Printf.printf \"abs(x): %s\\n\" (data_to_string (abs x));\n Printf.printf \"x²: %s\\n\" (data_to_string (square x));\n Printf.printf \"√|x|: %s\\n\" (data_to_string (sqrt (abs x)));\n Printf.printf \"exp(x): %s\\n\" (data_to_string (exp x));\n Printf.printf \"sign(x): %s\\n\\n\" (data_to_string (sign x));\n\n (* --- Clamp: cap sensor readings to a valid range --- *)\n let readings =\n create float64 [| 6 |] [| -5.0; 12.0; 105.0; 42.0; -1.0; 99.0 |]\n in\n let clamped = clamp ~min:0.0 ~max:100.0 readings in\n Printf.printf \"Sensor readings: %s\\n\" (data_to_string readings);\n Printf.printf \"Clamped [0,100]: %s\\n\" (data_to_string clamped)\n"
},
{
"path": "packages/nx/examples/03-indexing-and-slicing/README.md",
"content": "# `03-indexing-and-slicing`\n\nSelect, slice, and mask — extract exactly the data you need. This example uses\na grade book to demonstrate every way Nx lets you reach into an array.\n\n```bash\ndune exec nx/examples/03-indexing-and-slicing/main.exe\n```\n\n## What You'll Learn\n\n- Reading single elements with `item`\n- Selecting rows and columns with `I` and `A`\n- Range slicing with `R` and strided slicing with `Rs`\n- Infix indexing syntax: `.%{}` and `.${}`\n- Boolean masks with `compress` and `where`\n- Picking rows by index with `take`\n\n## Key Functions\n\n| Function / Index | Purpose |\n| ----------------------------- | -------------------------------------- |\n| `item [i; j] t` | Extract a single OCaml scalar |\n| `I n` | Select index `n` along one axis |\n| `A` | Select all indices along an axis |\n| `R (start, stop)` | Half-open range `[start, stop)` |\n| `Rs (start, stop, step)` | Range with stride |\n| `t.${[...]}` | Infix slicing (synonym for `slice`) |\n| `compress ~axis ~condition t` | Keep rows/cols where condition is true |\n| `where cond then_ else_` | Element-wise conditional selection |\n| `take ~axis indices t` | Gather rows by integer indices |\n| `greater_s t scalar` | Element-wise `t > scalar` → bool mask |\n\n## Output Walkthrough\n\nThe example starts with a 5×4 grade book (5 students, 4 subjects):\n\n```\nGrade book (students × subjects):\n[[88, 72, 95, 83],\n [45, 90, 67, 78],\n [92, 85, 91, 70],\n [76, 63, 80, 95],\n [60, 78, 55, 82]]\n```\n\n### Single element\n\n```ocaml\nitem [ 0; 1 ] grades (* → 72.0 *)\n```\n\n### Row and column selection\n\nThe infix `.${[...]}` operator makes slicing readable. `I n` picks one index,\n`A` keeps the full axis:\n\n```ocaml\ngrades.${[ I 2; A ]} (* student 2, all subjects → [92, 85, 91, 70] *)\ngrades.${[ A; I 0 ]} (* all students, Math → [88, 45, 92, 76, 60] *)\n```\n\n### Range and strided slicing\n\n`R (start, stop)` is a half-open range. `Rs (start, stop, step)` adds a stride:\n\n```ocaml\ngrades.${[ R (1, 4); R (0, 2) ]} (* students 1-3, Math & Science *)\ngrades.${[ Rs (0, 5, 2); Rs (0, 4, 2) ]} (* every other student & subject *)\n```\n\n### Boolean masks\n\nBuild a boolean mask, then use `compress` to filter rows:\n\n```ocaml\nlet high_math = greater_s (grades.${[ A; I 0 ]}) 85.0 in\ncompress ~axis:0 ~condition:high_math grades\n```\n\n```\nMath > 85 mask: [true, false, true, false, false]\nStudents with Math > 85:\n[[88, 72, 95, 83],\n [92, 85, 91, 70]]\n```\n\n### Conditional replacement\n\n`where` replaces elements based on a condition — here, flooring all grades\nbelow 60:\n\n```ocaml\nwhere (less_s grades 60.0) (full float64 [| 5; 4 |] 60.0) grades\n```\n\n## Index Types at a Glance\n\n| Index | Meaning | Example |\n| -------------- | -------------- | ------------------------------ |\n| `I n` | Single index | `I 2` — third element |\n| `A` | All indices | `A` — keep entire axis |\n| `R (a, b)` | Range `[a, b)` | `R (1, 4)` — indices 1, 2, 3 |\n| `Rs (a, b, s)` | Strided range | `Rs (0, 10, 2)` — even indices |\n| `L [...]` | Explicit list | `L [0; 3; 7]` — pick specific |\n| `M mask` | Boolean mask | `M bool_array` — where true |\n| `N` | New axis | `N` — insert dimension |\n\n## Try It\n\n1. Extract the Art column (column 3) for all students.\n2. Use `Rs (4, -1, -1)` to reverse the student order (negative step).\n3. Find students whose average grade across all subjects exceeds 80 using\n `mean ~axes:[1]` and a boolean mask.\n\n## Next Steps\n\nContinue to [04-reshaping-and-broadcasting](../04-reshaping-and-broadcasting/)\nto learn how to change array shapes and let broadcasting align dimensions\nautomatically.\n"
},
{
"path": "packages/nx/examples/03-indexing-and-slicing/dune",
"content": "(executable\n (name main)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/03-indexing-and-slicing/main.ml",
"content": "(** Select, slice, and mask — extract exactly the data you need.\n\n A grade book of 5 students across 4 subjects. We'll pull out individual\n scores, entire rows and columns, ranges, and use boolean masks to find top\n performers. *)\n\nopen Nx\nopen Nx.Infix\n\nlet () =\n (* Grade book: 5 students × 4 subjects (Math, Science, English, Art). *)\n let grades =\n create float64 [| 5; 4 |]\n [|\n 88.0;\n 72.0;\n 95.0;\n 83.0;\n 45.0;\n 90.0;\n 67.0;\n 78.0;\n 92.0;\n 85.0;\n 91.0;\n 70.0;\n 76.0;\n 63.0;\n 80.0;\n 95.0;\n 60.0;\n 78.0;\n 55.0;\n 82.0;\n |]\n in\n Printf.printf \"Grade book (students × subjects):\\n%s\\n\\n\"\n (data_to_string grades);\n\n (* Single element: student 0's Science score (row 0, col 1). *)\n let score = item [ 0; 1 ] grades in\n Printf.printf \"Student 0, Science: %.0f\\n\\n\" score;\n\n (* Entire row: all of student 2's grades. *)\n let student_2 = grades.${[ I 2; A ]} in\n Printf.printf \"Student 2 (all subjects): %s\\n\\n\" (data_to_string student_2);\n\n (* Entire column: everyone's Math scores (column 0). *)\n let math = grades.${[ A; I 0 ]} in\n Printf.printf \"Math scores (all students): %s\\n\\n\" (data_to_string math);\n\n (* Range: students 1-3, first two subjects. *)\n let subset = grades.${[ R (1, 4); R (0, 2) ]} in\n Printf.printf \"Students 1-3, Math & Science:\\n%s\\n\\n\" (data_to_string subset);\n\n (* Strided: every other student, every other subject. *)\n let strided = grades.${[ Rs (0, 5, 2); Rs (0, 4, 2) ]} in\n Printf.printf \"Every other student & subject:\\n%s\\n\\n\"\n (data_to_string strided);\n\n (* Boolean mask: which students scored above 85 in Math? *)\n let math_scores = grades.${[ A; I 0 ]} in\n let high_math = greater_s math_scores 85.0 in\n Printf.printf \"Math > 85 mask: %s\\n\" (data_to_string high_math);\n\n let top_students = compress ~axis:0 ~condition:high_math grades in\n Printf.printf \"Students with Math > 85:\\n%s\\n\\n\" (data_to_string top_students);\n\n (* where: replace failing grades (<60) with 60. *)\n let passing =\n where (less_s grades 60.0) (full float64 [| 5; 4 |] 60.0) grades\n in\n Printf.printf \"After floor at 60:\\n%s\\n\\n\" (data_to_string passing);\n\n (* take: select specific students by index. *)\n let picks = take ~axis:0 (create int32 [| 3 |] [| 0l; 2l; 4l |]) grades in\n Printf.printf \"Students 0, 2, 4:\\n%s\\n\" (data_to_string picks)\n"
},
{
"path": "packages/nx/examples/04-reshaping-and-broadcasting/README.md",
"content": "# `04-reshaping-and-broadcasting`\n\nChange array shapes and let broadcasting align dimensions automatically. This\nexample reshapes a flat signal into frames, centers data by subtracting column\nmeans, and builds an outer product — all without explicit loops.\n\n```bash\ndune exec nx/examples/04-reshaping-and-broadcasting/main.exe\n```\n\n## What You'll Learn\n\n- Reshaping flat arrays into multi-dimensional frames with `reshape`\n- Flattening back to 1D with `flatten`\n- Transposing rows and columns\n- Stacking arrays vertically and horizontally: `vstack`, `hstack`\n- Broadcasting: how `keepdims` enables operations on different-shaped arrays\n- Building outer products via broadcasting\n- Adding and removing dimensions with `expand_dims` and `squeeze`\n\n## Key Functions\n\n| Function | Purpose |\n| --------------------- | ------------------------------------------------------ |\n| `reshape shape t` | Change array shape (total elements must match) |\n| `flatten t` | Collapse all dimensions into 1D |\n| `transpose t` | Reverse all axes (swap rows and columns) |\n| `vstack ts` | Stack arrays vertically (along axis 0) |\n| `hstack ts` | Stack arrays horizontally (along axis 1) |\n| `expand_dims axes t` | Insert size-1 dimensions at specified positions |\n| `squeeze t` | Remove all size-1 dimensions |\n| `mean ~keepdims:true` | Reduce while keeping axis as size 1 (for broadcasting) |\n\n## Output Walkthrough\n\nReshape a flat 12-element signal into a 3×4 matrix of frames:\n\n```ocaml\nlet signal = arange_f float64 0.0 12.0 1.0 in\nlet frames = reshape [| 3; 4 |] signal\n```\n\n```\nFlat signal (12 samples):\n[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]\n\nReshaped into 3 frames of 4:\n[[0, 1, 2, 3],\n [4, 5, 6, 7],\n [8, 9, 10, 11]]\n```\n\n### Broadcasting in action\n\nSubtracting column means from data. The `keepdims:true` parameter gives\nthe mean shape `[1; 3]` instead of `[3]`, which broadcasts against `[4; 3]`:\n\n```ocaml\nlet col_means = mean ~axes:[ 0 ] ~keepdims:true data in\nlet centered = data - col_means\n```\n\n### Outer product via broadcasting\n\nReshape vectors into compatible shapes and multiply — no loops needed:\n\n```ocaml\nlet outer = reshape [| 4; 1 |] x * reshape [| 1; 3 |] y\n```\n\n```\nOuter product (x × y):\n[[10, 20, 30],\n [20, 40, 60],\n [30, 60, 90],\n [40, 80, 120]]\n```\n\n## Broadcasting Rules\n\nTwo dimensions are compatible for broadcasting when they are either:\n1. Equal, or\n2. One of them is 1\n\nWhen dimensions differ, the size-1 dimension is stretched to match. This is\nwhy `keepdims:true` is essential for reductions used in arithmetic.\n\n## Try It\n\n1. Reshape the signal into `[4; 3]` instead of `[3; 4]` and compare with the\n transpose of the original frames.\n2. Stack three 1D arrays of different values with `vstack`, then compute\n row-wise means using `mean ~axes:[1]`.\n3. Compute an outer product of two vectors of different lengths (e.g., 5 and 3)\n using `reshape` and broadcasting.\n\n## Next Steps\n\nContinue to [05-reductions-and-statistics](../05-reductions-and-statistics/) to\nlearn how to summarize data with aggregations along any axis.\n"
},
{
"path": "packages/nx/examples/04-reshaping-and-broadcasting/dune",
"content": "(executable\n (name main)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/04-reshaping-and-broadcasting/main.ml",
"content": "(** Change array shapes and let broadcasting align dimensions automatically.\n\n Reshape a flat signal into frames, center data by subtracting column means\n (broadcasting in action), and build an outer product without any loops. *)\n\nopen Nx\nopen Nx.Infix\n\nlet () =\n (* --- Reshape: flat signal → frames --- *)\n let signal = arange_f float64 0.0 12.0 1.0 in\n Printf.printf \"Flat signal (12 samples):\\n%s\\n\\n\" (data_to_string signal);\n\n let frames = reshape [| 3; 4 |] signal in\n Printf.printf \"Reshaped into 3 frames of 4:\\n%s\\n\\n\" (data_to_string frames);\n\n let flat_again = flatten frames in\n Printf.printf \"Flattened back: %s\\n\\n\" (data_to_string flat_again);\n\n (* --- Transpose: swap rows and columns --- *)\n Printf.printf \"Transposed:\\n%s\\n\\n\" (data_to_string (transpose frames));\n\n (* --- Stacking arrays --- *)\n let a = create float64 [| 3 |] [| 1.0; 2.0; 3.0 |] in\n let b = create float64 [| 3 |] [| 4.0; 5.0; 6.0 |] in\n Printf.printf \"vstack [a; b]:\\n%s\\n\" (data_to_string (vstack [ a; b ]));\n Printf.printf \"hstack [a; b]: %s\\n\\n\" (data_to_string (hstack [ a; b ]));\n\n (* --- Broadcasting: subtract column means to center data --- *)\n let data =\n create float64 [| 4; 3 |]\n [|\n 10.0;\n 200.0;\n 3000.0;\n 20.0;\n 400.0;\n 1000.0;\n 30.0;\n 100.0;\n 2000.0;\n 40.0;\n 300.0;\n 4000.0;\n |]\n in\n Printf.printf \"Raw data (4 samples × 3 features):\\n%s\\n\" (data_to_string data);\n\n (* Mean along axis 0 with keepdims — shape [1; 3] broadcasts against [4;\n 3]. *)\n let col_means = mean ~axes:[ 0 ] ~keepdims:true data in\n Printf.printf \"Column means: %s\\n\" (data_to_string col_means);\n\n let centered = data - col_means in\n Printf.printf \"Centered (zero-mean columns):\\n%s\\n\\n\"\n (data_to_string centered);\n\n (* --- Outer product via broadcasting --- *)\n let x = create float64 [| 4 |] [| 1.0; 2.0; 3.0; 4.0 |] in\n let y = create float64 [| 3 |] [| 10.0; 20.0; 30.0 |] in\n\n (* x as column [4;1], y as row [1;3] → result is [4;3]. *)\n let outer = reshape [| 4; 1 |] x * reshape [| 1; 3 |] y in\n Printf.printf \"x = %s\\n\" (data_to_string x);\n Printf.printf \"y = %s\\n\" (data_to_string y);\n Printf.printf \"Outer product (x × y):\\n%s\\n\\n\" (data_to_string outer);\n\n (* --- expand_dims / squeeze --- *)\n let v = arange float64 0 4 1 in\n let row = expand_dims [ 0 ] v in\n let col = expand_dims [ 1 ] v in\n Printf.printf \"Vector: shape %s → %s\\n\"\n (shape_to_string (shape v))\n (data_to_string v);\n Printf.printf \"Row vector: shape %s → %s\\n\"\n (shape_to_string (shape row))\n (data_to_string row);\n Printf.printf \"Col vector: shape %s\\n%s\\n\"\n (shape_to_string (shape col))\n (data_to_string col)\n"
},
{
"path": "packages/nx/examples/05-reductions-and-statistics/README.md",
"content": "# `05-reductions-and-statistics`\n\nSummarize data with reductions — means, variances, and aggregations along any\naxis. This example analyzes daily temperature readings across four cities.\n\n```bash\ndune exec nx/examples/05-reductions-and-statistics/main.exe\n```\n\n## What You'll Learn\n\n- Reducing along specific axes with `mean`, `std`, `sum`\n- Finding extremes and their positions with `min`, `max`, `argmax`\n- Computing running totals with `cumsum`\n- Preserving dimensions for broadcasting with `keepdims`\n- Detecting outliers using z-score normalization\n- Testing conditions with `all` and `any`\n\n## Key Functions\n\n| Function | Purpose |\n| ----------------------------- | ------------------------------------------- |\n| `mean ~axes t` | Average values along specified axes |\n| `std ~axes t` | Standard deviation along axes |\n| `min t` / `max t` | Global minimum / maximum |\n| `min ~axes t` / `max ~axes t` | Per-axis minimum / maximum |\n| `argmax ~axis t` | Index of the maximum along an axis |\n| `cumsum ~axis t` | Cumulative sum along an axis |\n| `all t` / `any t` | Test if all / any elements are true |\n| `greater_s t s` | Element-wise `t > s` returning a bool array |\n| `less_s t s` | Element-wise `t < s` returning a bool array |\n\n## Output Walkthrough\n\nThe dataset is a 4×7 matrix — 4 cities, 7 days of temperature readings:\n\n```ocaml\nlet city_means = mean ~axes:[ 1 ] temps in\n```\n\n```\nCity averages:\n Paris mean=22.9 std=2.3\n Cairo mean=32.0 std=2.1\n Helsinki mean=-5.6 std=2.6\n London mean=14.9 std=1.3\n```\n\n### Axis semantics\n\n- `~axes:[1]` reduces across columns (days) → one value per city\n- `~axes:[0]` reduces across rows (cities) → one value per day\n- No axis → reduces everything to a scalar\n\n### Outlier detection with z-scores\n\nUsing `keepdims:true` to broadcast the mean and std against the original data:\n\n```ocaml\nlet mu = mean ~axes:[ 1 ] ~keepdims:true temps in\nlet sigma = std ~axes:[ 1 ] ~keepdims:true temps in\nlet z_scores = (temps - mu) / sigma in\nlet outlier_mask = greater_s (abs z_scores) 1.5\n```\n\n### Condition testing\n\n```ocaml\nlet all_above_zero = all (greater_s temps 0.0) in (* false — Helsinki *)\nlet any_below_neg5 = any (less_s temps (-5.0)) in (* true — Helsinki *)\n```\n\n## Try It\n\n1. Compute the daily average across all cities with `mean ~axes:[0]` and find\n which day was warmest on average.\n2. Use `cumsum ~axis:1` on the full temperature matrix to see running totals\n per city.\n3. Find the day with the smallest temperature range across cities using\n `max ~axes:[0]` minus `min ~axes:[0]`.\n\n## Next Steps\n\nContinue to [06-random-numbers](../06-random-numbers/) to generate synthetic\ndata with controlled, reproducible distributions.\n"
},
{
"path": "packages/nx/examples/05-reductions-and-statistics/dune",
"content": "(executable\n (name main)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/05-reductions-and-statistics/main.ml",
"content": "(** Summarize data with reductions — means, variances, and aggregations along\n any axis.\n\n Analyze daily temperature readings across four cities. Compute averages,\n find extremes, track running totals, and flag outliers. *)\n\nopen Nx\nopen Nx.Infix\n\nlet () =\n (* Daily temperatures (°C) for 4 cities over 7 days. Rows = cities, columns =\n days. *)\n let temps =\n create float64 [| 4; 7 |]\n [|\n 22.0;\n 24.0;\n 19.0;\n 25.0;\n 23.0;\n 21.0;\n 26.0;\n (* Paris *)\n 30.0;\n 32.0;\n 35.0;\n 31.0;\n 29.0;\n 33.0;\n 34.0;\n (* Cairo *)\n -5.0;\n -8.0;\n -3.0;\n -10.0;\n -2.0;\n -7.0;\n -4.0;\n (* Helsinki *)\n 15.0;\n 14.0;\n 16.0;\n 13.0;\n 17.0;\n 15.0;\n 14.0;\n (* London *)\n |]\n in\n let cities = [| \"Paris\"; \"Cairo\"; \"Helsinki\"; \"London\" |] in\n Printf.printf \"Daily temperatures (4 cities × 7 days):\\n%s\\n\\n\"\n (data_to_string temps);\n\n (* --- Per-city statistics (reduce along axis 1 = across days) --- *)\n let city_means = mean ~axes:[ 1 ] temps in\n let city_stds = std ~axes:[ 1 ] temps in\n Printf.printf \"City averages:\\n\";\n for i = 0 to 3 do\n Printf.printf \" %-10s mean=%.1f std=%.1f\\n\" cities.(i)\n (item [ i ] city_means) (item [ i ] city_stds)\n done;\n print_newline ();\n\n (* --- Hottest day per city (argmax along axis 1) --- *)\n let hottest_day = argmax ~axis:1 temps in\n Printf.printf \"Hottest day per city:\\n\";\n for i = 0 to 3 do\n Printf.printf \" %-10s day %ld\\n\" cities.(i) (item [ i ] hottest_day)\n done;\n print_newline ();\n\n (* --- Global extremes --- *)\n Printf.printf \"Warmest reading: %.1f°C\\n\" (item [] (max temps));\n Printf.printf \"Coldest reading: %.1f°C\\n\\n\" (item [] (min temps));\n\n (* --- Cumulative sum: running total of Cairo's temperatures --- *)\n let cairo = temps.${[ I 1; A ]} in\n let cumulative = cumsum ~axis:0 cairo in\n Printf.printf \"Cairo daily: %s\\n\" (data_to_string cairo);\n Printf.printf \"Cairo cumulative: %s\\n\\n\" (data_to_string cumulative);\n\n (* --- Outlier detection with z-scores --- *)\n let mu = mean ~axes:[ 1 ] ~keepdims:true temps in\n let sigma = std ~axes:[ 1 ] ~keepdims:true temps in\n let z_scores = (temps - mu) / sigma in\n let outlier_mask = greater_s (abs z_scores) 1.5 in\n Printf.printf \"Z-scores:\\n%s\\n\" (data_to_string z_scores);\n Printf.printf \"Outliers (|z| > 1.5): %s\\n\\n\" (data_to_string outlier_mask);\n\n (* --- Check if all/any values meet a condition --- *)\n let all_above_zero = all (greater_s temps 0.0) in\n let any_below_neg5 = any (less_s temps (-5.0)) in\n Printf.printf \"All temps > 0? %b\\n\" (item [] all_above_zero);\n Printf.printf \"Any temp < -5? %b\\n\" (item [] any_below_neg5)\n"
},
{
"path": "packages/nx/examples/06-random-numbers/README.md",
"content": "# `06-random-numbers`\n\nImplicit RNG with reproducible scopes — generate distributions, sample,\nand shuffle. Wrap code in `Rng.run ~seed` for deterministic results.\n\n```bash\ndune exec nx/examples/06-random-numbers/main.exe\n```\n\n## What You'll Learn\n\n- Generating uniform, normal, and integer distributions\n- Running a Monte Carlo simulation to estimate pi\n- Creating synthetic training data with controlled noise\n- Verifying reproducibility with `Rng.run ~seed`\n- Shuffling arrays with `Rng.shuffle`\n\n## Key Functions\n\n| Function | Purpose |\n| ---------------------------------------- | ------------------------------------------------- |\n| `Rng.run ~seed f` | Execute `f` in a deterministic RNG scope |\n| `Rng.uniform dtype shape` | Uniform random values in [0, 1) |\n| `Rng.normal dtype shape` | Standard normal distribution (mean=0, std=1) |\n| `Rng.randint ~high dtype shape low` | Random integers in [low, high) |\n| `Rng.shuffle t` | Randomly permute array elements |\n| `rand dtype shape` | Shorthand for uniform random values |\n\n## Output Walkthrough\n\n### Monte Carlo pi estimation\n\nDrop random points in a unit square. The fraction inside the unit circle\napproximates pi/4:\n\n```ocaml\nlet xs = rand float64 [| n |] in\nlet ys = rand float64 [| n |] in\nlet inside = less_s ((xs * xs) + (ys * ys)) 1.0 in\nlet pi_est = item [] (sum (cast Float64 inside)) *. 4.0 /. Float.of_int n\n```\n\n```\nMonte Carlo pi (100000 points): 3.1420 (actual: 3.1416)\n```\n\n### Reproducibility\n\nSame seed always produces the same numbers:\n\n```ocaml\nlet a = Rng.run ~seed:99 (fun () -> Rng.normal Float64 [| 3 |]) in\nlet b = Rng.run ~seed:99 (fun () -> Rng.normal Float64 [| 3 |]) in\n(* Identical? true *)\n```\n\n## Try It\n\n1. Roll 1000 dice with `Rng.randint` and compute the mean — it should\n approach the theoretical expected value of 3.5.\n2. Increase the Monte Carlo sample count to 1,000,000 and observe how the pi\n estimate improves.\n3. Generate two clusters of 2D points (one centered at origin, one at (3, 3))\n using `Rng.normal` with offsets.\n\n## Next Steps\n\nContinue to [07-linear-algebra](../07-linear-algebra/) to learn matrix\noperations, decompositions, and solving linear systems.\n"
},
{
"path": "packages/nx/examples/06-random-numbers/dune",
"content": "(executable\n (name main)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/06-random-numbers/main.ml",
"content": "(** Implicit RNG with reproducible scopes — generate distributions, sample, and\n shuffle.\n\n Roll dice, estimate pi with Monte Carlo, and generate noisy training data.\n Every result is reproducible inside an [Rng.run] scope: same seed, same\n numbers. Outside any scope the global fallback provides convenient but\n non-reproducible randomness. *)\n\nopen Nx\nopen Nx.Infix\n\nlet () =\n (* --- Dice simulation: roll 10 six-sided dice --- *)\n let dice = randint Int32 ~high:7 [| 10 |] 1 in\n Printf.printf \"10 dice rolls: %s\\n\\n\" (data_to_string dice);\n\n (* --- Uniform random floats in [0, 1) --- *)\n let uniform = rand Float64 [| 5 |] in\n Printf.printf \"Uniform [0,1): %s\\n\\n\" (data_to_string uniform);\n\n (* --- Normal distribution (mean=0, std=1) --- *)\n let normal = randn Float64 [| 5 |] in\n Printf.printf \"Normal(0,1): %s\\n\\n\" (data_to_string normal);\n\n (* --- Monte Carlo pi estimation ---\n\n Drop N random points in a unit square. The fraction landing inside the unit\n circle (distance from origin < 1) approximates pi/4. *)\n let n = 100_000 in\n let xs = rand float64 [| n |] in\n let ys = rand float64 [| n |] in\n let inside = less_s ((xs * xs) + (ys * ys)) 1.0 in\n let count = sum (cast Float64 inside) in\n let pi_est = item [] count *. 4.0 /. Float.of_int n in\n Printf.printf \"Monte Carlo pi (%d points): %.4f (actual: %.4f)\\n\\n\" n pi_est\n Float.pi;\n\n (* --- Synthetic training data: y = 3x + 2 + noise --- *)\n let x = rand Float64 [| 8 |] in\n let noise = randn Float64 [| 8 |] *$ 0.1 in\n let y = (x *$ 3.0) +$ 2.0 + noise in\n Printf.printf \"x: %s\\n\" (data_to_string x);\n Printf.printf \"y ~ 3x+2: %s\\n\\n\" (data_to_string y);\n\n (* --- Reproducibility: Rng.run ~seed gives the same result --- *)\n let a = Rng.run ~seed:99 (fun () -> randn Float64 [| 3 |]) in\n let b = Rng.run ~seed:99 (fun () -> randn Float64 [| 3 |]) in\n Printf.printf \"Same seed, run 1: %s\\n\" (data_to_string a);\n Printf.printf \"Same seed, run 2: %s\\n\" (data_to_string b);\n Printf.printf \"Identical? %b\\n\\n\" (item [] (all (equal a b)));\n\n (* --- Shuffle: random permutation of a dataset --- *)\n let data = arange int32 0 8 1 in\n let shuffled = shuffle data in\n Printf.printf \"Original: %s\\n\" (data_to_string data);\n Printf.printf \"Shuffled: %s\\n\" (data_to_string shuffled)\n"
},
{
"path": "packages/nx/examples/07-linear-algebra/README.md",
"content": "# `07-linear-algebra`\n\nSolve systems, decompose matrices, and fit models — linear algebra made\npractical. This example covers matrix multiplication, linear solves, least\nsquares fitting, eigendecomposition, and SVD.\n\n```bash\ndune exec nx/examples/07-linear-algebra/main.exe\n```\n\n## What You'll Learn\n\n- Matrix multiplication with `@@` and dot products with `<.>`\n- Solving linear systems with `/@`\n- Computing inverses, determinants, and norms\n- Fitting a line to data with least squares (`lstsq`)\n- Eigendecomposition of symmetric matrices (`eigh`)\n- Singular value decomposition and reconstruction (`svd`)\n\n## Key Functions\n\n| Function | Purpose |\n| ----------- | -------------------------------------------------- |\n| `a @@ b` | Matrix multiplication |\n| `u <.> v` | Vector dot product |\n| `a /@ b` | Solve linear system Ax = b |\n| `inv m` | Matrix inverse |\n| `det m` | Determinant |\n| `norm m` | Matrix norm (Frobenius by default) |\n| `lstsq a b` | Least squares solution to overdetermined system |\n| `eigh m` | Eigenvalues and eigenvectors of a symmetric matrix |\n| `svd m` | Singular value decomposition (U, S, Vt) |\n| `diag v` | Create diagonal matrix from a vector |\n\n## Output Walkthrough\n\n### Matrix multiplication\n\n```ocaml\nlet a = create float64 [| 2; 3 |] [| 1.; 2.; 3.; 4.; 5.; 6. |] in\nlet b = create float64 [| 3; 2 |] [| 7.; 8.; 9.; 10.; 11.; 12. |] in\na @@ b (* [2; 2] result *)\n```\n\n### Solving linear systems\n\nThe `/@` operator solves Ax = b for x:\n\n```ocaml\nlet x = coeff /@ rhs (* x = [2; 3; -1] *)\n```\n\n### Inverse verification\n\n```ocaml\nlet m_inv = inv m in\nm @@ m_inv (* ≈ identity matrix *)\n```\n\n### Least squares fitting\n\nBuild a design matrix [x, 1] and solve for slope and intercept:\n\n```ocaml\nlet design = hstack [ x_col; ones float64 [| 6; 1 |] ] in\nlet coeffs, _, _, _ = lstsq design y_col in\n(* m ≈ 1.97, c ≈ 1.03 *)\n```\n\n### SVD decomposition and reconstruction\n\n```ocaml\nlet u_mat, s_vec, vt = svd data in\nlet reconstructed = u_mat.${[ A; R (0, 2) ]} @@ diag s_vec @@ vt\n(* reconstructed ≈ original *)\n```\n\n## Try It\n\n1. Solve a different 3×3 system and verify the solution by computing\n `coeff @@ x` — it should match the right-hand side.\n2. Extend the least squares example to fit a quadratic by adding an x^2\n column to the design matrix.\n3. Use SVD for low-rank approximation: zero out the smallest singular value,\n reconstruct, and compare to the original.\n\n## Next Steps\n\nContinue to [08-signal-processing](../08-signal-processing/) to apply\nfrequency analysis with FFT.\n"
},
{
"path": "packages/nx/examples/07-linear-algebra/dune",
"content": "(executable\n (name main)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/07-linear-algebra/main.ml",
"content": "(** Solve systems, decompose matrices, and fit models — linear algebra made\n practical.\n\n Fit a line to noisy data with least squares, verify matrix inverses,\n decompose matrices with SVD and eigendecomposition. *)\n\nopen Nx\nopen Nx.Infix\n\nlet () =\n (* --- Matrix multiplication --- *)\n let a = create float64 [| 2; 3 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0 |] in\n let b = create float64 [| 3; 2 |] [| 7.0; 8.0; 9.0; 10.0; 11.0; 12.0 |] in\n Printf.printf \"A (2×3):\\n%s\\n\" (data_to_string a);\n Printf.printf \"B (3×2):\\n%s\\n\" (data_to_string b);\n Printf.printf \"A @@ B:\\n%s\\n\\n\" (data_to_string (a @@ b));\n\n (* --- Dot product --- *)\n let u = create float64 [| 3 |] [| 1.0; 2.0; 3.0 |] in\n let v = create float64 [| 3 |] [| 4.0; 5.0; 6.0 |] in\n Printf.printf \"u · v = %s\\n\\n\" (data_to_string (u <.> v));\n\n (* --- Solving linear systems: A x = b --- *)\n let coeff =\n create float64 [| 3; 3 |]\n [| 2.0; 1.0; -1.0; -3.0; -1.0; 2.0; -2.0; 1.0; 2.0 |]\n in\n let rhs = create float64 [| 3; 1 |] [| 8.0; -11.0; -3.0 |] in\n let x = coeff /@ rhs in\n Printf.printf \"System Ax = b:\\n\";\n Printf.printf \"A:\\n%s\\n\" (data_to_string coeff);\n Printf.printf \"b: %s\\n\" (data_to_string (flatten rhs));\n Printf.printf \"x: %s\\n\\n\" (data_to_string (flatten x));\n\n (* --- Inverse: verify A @@ inv(A) ≈ I --- *)\n let m = create float64 [| 2; 2 |] [| 4.0; 7.0; 2.0; 6.0 |] in\n let m_inv = inv m in\n let product = m @@ m_inv in\n Printf.printf \"M:\\n%s\\n\" (data_to_string m);\n Printf.printf \"M⁻¹:\\n%s\\n\" (data_to_string m_inv);\n Printf.printf \"M × M⁻¹ ≈ I:\\n%s\\n\\n\" (data_to_string product);\n\n (* --- Determinant and norm --- *)\n Printf.printf \"det(M) = %.1f\\n\" (item [] (det m));\n Printf.printf \"‖M‖_F = %.4f\\n\\n\" (item [] (norm m));\n\n (* --- Least squares: fit y = mx + c to noisy data --- *)\n let x_data = create float64 [| 6 |] [| 0.0; 1.0; 2.0; 3.0; 4.0; 5.0 |] in\n let y_data = create float64 [| 6 |] [| 1.1; 2.9; 5.2; 6.8; 9.1; 10.8 |] in\n\n (* Build design matrix: [x, 1] for y = m*x + c *)\n let x_col = reshape [| 6; 1 |] x_data in\n let design = hstack [ x_col; ones float64 [| 6; 1 |] ] in\n let y_col = reshape [| 6; 1 |] y_data in\n\n let coeffs, _residuals, _rank, _sv = lstsq design y_col in\n Printf.printf \"Least squares fit y = m·x + c:\\n\";\n Printf.printf \" m = %.4f\\n\" (item [ 0; 0 ] coeffs);\n Printf.printf \" c = %.4f\\n\\n\" (item [ 1; 0 ] coeffs);\n\n (* --- Eigendecomposition of a symmetric matrix --- *)\n let sym =\n create float64 [| 3; 3 |]\n [| 2.0; -1.0; 0.0; -1.0; 2.0; -1.0; 0.0; -1.0; 2.0 |]\n in\n let eigenvalues, eigenvectors = eigh sym in\n Printf.printf \"Symmetric matrix:\\n%s\\n\" (data_to_string sym);\n Printf.printf \"Eigenvalues: %s\\n\" (data_to_string eigenvalues);\n Printf.printf \"Eigenvectors:\\n%s\\n\\n\" (data_to_string eigenvectors);\n\n (* --- SVD: decompose and reconstruct --- *)\n let data = create float64 [| 3; 2 |] [| 1.0; 2.0; 3.0; 4.0; 5.0; 6.0 |] in\n let u_mat, s_vec, vt = svd data in\n (* Reconstruct: U @ diag(S) @ Vt *)\n let s_diag = diag s_vec in\n let reconstructed = u_mat.${[ A; R (0, 2) ]} @@ s_diag @@ vt in\n Printf.printf \"Original:\\n%s\\n\" (data_to_string data);\n Printf.printf \"Singular values: %s\\n\" (data_to_string s_vec);\n Printf.printf \"Reconstructed (U·S·Vt):\\n%s\\n\" (data_to_string reconstructed)\n"
},
{
"path": "packages/nx/examples/08-signal-processing/README.md",
"content": "# `08-signal-processing`\n\nAnalyze frequencies with FFT — decompose signals and filter noise. This example\nbuilds a signal from two sine waves plus noise, identifies the component\nfrequencies, and filters the noise in the frequency domain.\n\n```bash\ndune exec nx/examples/08-signal-processing/main.exe\n```\n\n## What You'll Learn\n\n- Constructing synthetic signals from sine waves and noise\n- Transforming to the frequency domain with `rfft`\n- Mapping frequency bins to Hz with `rfftfreq`\n- Identifying dominant frequency components by magnitude\n- Filtering noise by zeroing small-magnitude frequency bins\n- Reconstructing a clean signal with `irfft`\n\n## Key Functions\n\n| Function | Purpose |\n| ----------------------------- | ------------------------------------------------- |\n| `rfft t` | Real-valued FFT (time domain to frequency domain) |\n| `irfft ~n t` | Inverse real FFT (frequency domain back to time) |\n| `rfftfreq ~d n` | Frequency bin labels for `rfft` output |\n| `linspace dtype start stop n` | Evenly spaced time samples |\n| `sin t` | Element-wise sine |\n| `Rng.normal ~key dtype shape` | Gaussian noise |\n| `Nx.Infix` (`+`, `*`, `*$`) | Clean arithmetic on arrays |\n\n## Output Walkthrough\n\n### Signal construction\n\nA 256-sample signal at 256 Hz composed of two sine waves plus noise:\n\n```\nSignal: 256 samples at 256 Hz\nComponents: 5 Hz (amplitude 1.0) + 20 Hz (amplitude 0.5) + noise\n```\n\n### Frequency analysis\n\nAfter `rfft`, compute magnitudes scaled by 2/N to get amplitudes:\n\n```\nDominant frequencies:\n 5.0 Hz (magnitude 1.002)\n 20.0 Hz (magnitude 0.501)\n```\n\nThe FFT correctly recovers both sine components. Noise spreads across many\nbins with small magnitudes.\n\n### Noise filtering\n\nZero all frequency bins below a threshold, then reconstruct with `irfft`:\n\n```\nAfter filtering (threshold=0.2):\n Original first 8: [1.29, 1.16, 0.81, 0.44, ...]\n Filtered first 8: [1.00, 1.16, 0.79, 0.36, ...]\n```\n\nThe filtered signal retains the two sine waves while removing noise.\n\n### Frequency bins\n\nFor an N-sample real signal, `rfft` produces N/2 + 1 complex bins from 0 Hz\nto the Nyquist frequency (sample_rate / 2):\n\n```\nTotal bins: 129 (for 256-sample signal)\n```\n\n## Try It\n\n1. Add a third sine wave at 40 Hz and verify it appears in the dominant\n frequencies.\n2. Raise the filter threshold to 0.5 and observe how the 20 Hz component\n gets removed (its amplitude is only 0.5).\n3. Double the sample rate to 512 Hz and check how the frequency resolution\n changes.\n\n## Next Steps\n\nContinue to [09-image-processing](../09-image-processing/) to apply\nconvolution and pooling to 2D image data.\n"
},
{
"path": "packages/nx/examples/08-signal-processing/dune",
"content": "(executable\n (name main)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/08-signal-processing/main.ml",
"content": "(** Analyze frequencies with FFT — decompose signals and filter noise.\n\n Build a signal from two sine waves plus noise. Use the real FFT to identify\n component frequencies, then filter the noise and reconstruct a clean signal.\n*)\n\nopen Nx\nopen Nx.Infix\n\nlet () =\n (* Signal parameters. *)\n let n = 256 in\n let sample_rate = 256.0 in\n let dt = 1.0 /. sample_rate in\n\n (* Time axis: n samples at the given rate. *)\n let t = linspace float64 0.0 (Float.of_int n *. dt) n ~endpoint:false in\n\n (* Build signal: 5 Hz sine + 20 Hz sine + noise. *)\n let noise = randn Float64 [| n |] *$ 0.3 in\n let pi2 = 2.0 *. Float.pi in\n let signal_5hz = sin (t *$ (pi2 *. 5.0)) in\n let signal_20hz = sin (t *$ (pi2 *. 20.0)) *$ 0.5 in\n let signal = signal_5hz + signal_20hz + noise in\n\n Printf.printf \"Signal: %d samples at %.0f Hz\\n\" n sample_rate;\n Printf.printf\n \"Components: 5 Hz (amplitude 1.0) + 20 Hz (amplitude 0.5) + noise\\n\\n\";\n\n (* Show first 8 samples. *)\n Printf.printf \"First 8 samples: %s\\n\\n\"\n (data_to_string (slice [ R (0, 8) ] signal));\n\n (* --- Real FFT: transform to frequency domain --- *)\n let spectrum = rfft signal in\n let freqs = rfftfreq ~d:dt n in\n\n (* Magnitudes (scaled by 2/N for single-sided spectrum). Extract real and\n imaginary parts to compute |z| = sqrt(re² + im²). *)\n let spectrum_arr = to_array spectrum in\n let re =\n create float64 (shape spectrum)\n (Array.map (fun c -> c.Complex.re) spectrum_arr)\n in\n let im =\n create float64 (shape spectrum)\n (Array.map (fun c -> c.Complex.im) spectrum_arr)\n in\n let magnitudes = sqrt ((re * re) + (im * im)) *$ (2.0 /. Float.of_int n) in\n\n (* Find the dominant frequencies (magnitude > 0.3). *)\n Printf.printf \"Dominant frequencies:\\n\";\n let n_freqs = (shape magnitudes).(0) in\n for i = 0 to pred n_freqs do\n let mag = item [ i ] magnitudes in\n if Stdlib.( > ) mag 0.3 then\n Printf.printf \" %.1f Hz (magnitude %.3f)\\n\" (item [ i ] freqs) mag\n done;\n print_newline ();\n\n (* --- Filter: zero out small frequency components --- *)\n let threshold = 0.2 in\n let mag_arr = to_array magnitudes in\n let filtered =\n Array.mapi\n (fun i c ->\n if Stdlib.( < ) mag_arr.(i) threshold then Complex.zero else c)\n (to_array spectrum)\n in\n let clean_spectrum = create Complex128 (shape spectrum) filtered in\n\n (* Inverse FFT back to time domain. *)\n let clean_signal = irfft ~n clean_spectrum in\n\n Printf.printf \"After filtering (threshold=%.1f):\\n\" threshold;\n Printf.printf \" Original first 8: %s\\n\"\n (data_to_string (slice [ R (0, 8) ] signal));\n Printf.printf \" Filtered first 8: %s\\n\\n\"\n (data_to_string (slice [ R (0, 8) ] clean_signal));\n\n (* --- Frequency bins explained --- *)\n Printf.printf \"Frequency bins (first 10): %s\\n\"\n (data_to_string (slice [ R (0, 10) ] freqs));\n Printf.printf \"Total bins: %d (for %d-sample signal)\\n\" (shape freqs).(0) n\n"
},
{
"path": "packages/nx/examples/09-image-processing/README.md",
"content": "# `09-image-processing`\n\nLoad, transform, and save images as arrays — convolutions, pooling, and pixel\nmath. This example creates a synthetic grayscale image, blurs it, detects edges\nwith Sobel filters, and downsamples with max pooling.\n\n```bash\ndune exec nx/examples/09-image-processing/main.exe\n```\n\n## What You'll Learn\n\n- Creating synthetic images with `init` and pixel math\n- Applying 2D convolution with `correlate2d` (NCHW format)\n- Gaussian blur with a 3x3 kernel\n- Sobel edge detection (horizontal + vertical gradients)\n- Downsampling with `max_pool2d`\n- Converting between `UInt8` and `Float32` for computation\n- Saving arrays as PNG files with `Nx_io.save_image`\n\n## Key Functions\n\n| Function | Purpose |\n| --------------------------------------------- | -------------------------------------------------- |\n| `init UInt8 shape f` | Create an image by computing each pixel |\n| `correlate2d ~padding_mode:\\`Same img kernel` | 2D convolution (expects NCHW) |\n| `max_pool2d ~kernel_size ~stride img` | Downsample by taking max in each window |\n| `cast Float32 t` | Convert dtype for floating-point operations |\n| `clamp ~min ~max t` | Clamp values to a valid pixel range |\n| `contiguous t` | Ensure contiguous memory layout (required for I/O) |\n| `Nx_io.save_image path t` | Save a 2D (HxW) array as a grayscale PNG |\n\n## Output Walkthrough\n\n### Synthetic image\n\nA 64x64 horizontal gradient with a bright rectangle in the center:\n\n```ocaml\nlet img = init UInt8 [| h; w |] (fun idx ->\n let y = idx.(0) and x = idx.(1) in\n let base = x * 255 / (w - 1) in\n if y >= 16 && y < 48 && x >= 16 && x < 48 then 220 else base)\n```\n\n### NCHW format\n\nConvolution operations expect 4D tensors in NCHW format (batch, channels,\nheight, width). Convert with:\n\n```ocaml\nlet img_f = cast Float32 img |> contiguous |> reshape [| 1; 1; h; w |]\n```\n\n### Gaussian blur\n\nA 3x3 kernel with weights summing to 1, giving more weight to the center:\n\n```ocaml\nlet blur_kernel = create Float32 [| 1; 1; 3; 3 |]\n [| 1./16.; 2./16.; 1./16.;\n 2./16.; 4./16.; 2./16.;\n 1./16.; 2./16.; 1./16. |]\n```\n\n### Sobel edge detection\n\nCombines horizontal and vertical gradient magnitudes:\n\n```ocaml\nlet gx = correlate2d ~padding_mode:`Same img_f sobel_x in\nlet gy = correlate2d ~padding_mode:`Same img_f sobel_y in\nlet edges = sqrt (add (mul gx gx) (mul gy gy))\n```\n\n### Max pooling\n\n2x downsampling by taking the maximum in each 2x2 window:\n\n```\nSaved: pooled.png (64x64 -> 32x32)\n```\n\n## Output Files\n\nRunning this example creates four PNG files in the current directory:\n\n| File | Description |\n| -------------- | ------------------------------ |\n| `gradient.png` | Original synthetic image |\n| `blurred.png` | After Gaussian blur |\n| `edges.png` | Sobel edge detection result |\n| `pooled.png` | 2x downsampled via max pooling |\n\n## Try It\n\n1. Replace the blur kernel with a sharpening kernel:\n `[| 0.; -1.; 0.; -1.; 5.; -1.; 0.; -1.; 0. |]`\n2. Try a larger pooling window (`4, 4`) and observe the effect on image size\n and detail.\n3. Chain blur and edge detection: blur first to reduce noise, then apply Sobel.\n\n## Next Steps\n\nYou've completed the Nx examples! For machine learning workflows, see the\n[kaun examples](/docs/kaun/).\n"
},
{
"path": "packages/nx/examples/09-image-processing/dune",
"content": "(executable\n (name main)\n (libraries nx nx.io))\n"
},
{
"path": "packages/nx/examples/09-image-processing/main.ml",
"content": "(** Load, transform, and save images as arrays — convolutions, pooling, and\n pixel math.\n\n Create a synthetic grayscale gradient, blur it, detect edges with Sobel\n filters, and downsample with max pooling. Results are saved as PNG files. *)\n\nopen Nx\n\nlet () =\n let h = 64 and w = 64 in\n\n (* --- Create a gradient image with a bright rectangle --- *)\n let img =\n init UInt8 [| h; w |] (fun idx ->\n let y = idx.(0) and x = idx.(1) in\n (* Background: horizontal gradient. *)\n let base = x * 255 / (w - 1) in\n (* Bright rectangle in the center. *)\n if y >= 16 && y < 48 && x >= 16 && x < 48 then 220 else base)\n in\n Printf.printf \"Created %dx%d grayscale image\\n\" h w;\n\n (* Save the original. *)\n Nx_io.save_image \"gradient.png\" (contiguous img);\n Printf.printf \"Saved: gradient.png\\n\";\n\n (* --- Gaussian blur with a 3x3 kernel --- *)\n\n (* Convert to float for convolution. The scipy-style correlate works on raw\n spatial dims, so we use [H; W] directly. *)\n let img_f = cast Float32 img |> contiguous in\n\n let blur_kernel =\n create Float32 [| 3; 3 |]\n [|\n 1.0 /. 16.0;\n 2.0 /. 16.0;\n 1.0 /. 16.0;\n 2.0 /. 16.0;\n 4.0 /. 16.0;\n 2.0 /. 16.0;\n 1.0 /. 16.0;\n 2.0 /. 16.0;\n 1.0 /. 16.0;\n |]\n in\n let blurred = correlate ~padding:`Same img_f blur_kernel in\n let blurred_img =\n clamp ~min:0.0 ~max:255.0 blurred |> cast UInt8 |> contiguous\n in\n Nx_io.save_image \"blurred.png\" blurred_img;\n Printf.printf \"Saved: blurred.png\\n\";\n\n (* --- Sobel edge detection --- *)\n let sobel_x =\n create Float32 [| 3; 3 |]\n [| -1.0; 0.0; 1.0; -2.0; 0.0; 2.0; -1.0; 0.0; 1.0 |]\n in\n let sobel_y =\n create Float32 [| 3; 3 |]\n [| -1.0; -2.0; -1.0; 0.0; 0.0; 0.0; 1.0; 2.0; 1.0 |]\n in\n\n let gx = correlate ~padding:`Same img_f sobel_x in\n let gy = correlate ~padding:`Same img_f sobel_y in\n let edges = sqrt (add (mul gx gx) (mul gy gy)) in\n let edges_img = clamp ~min:0.0 ~max:255.0 edges |> cast UInt8 |> contiguous in\n Nx_io.save_image \"edges.png\" edges_img;\n Printf.printf \"Saved: edges.png\\n\";\n\n (* --- Max pooling: 2x downsample using maximum_filter --- *)\n let pooled =\n maximum_filter ~kernel_size:[| 2; 2 |] ~stride:[| 2; 2 |] img_f\n in\n let pool_h = (shape pooled).(0) and pool_w = (shape pooled).(1) in\n let pooled_img =\n clamp ~min:0.0 ~max:255.0 pooled |> cast UInt8 |> contiguous\n in\n Nx_io.save_image \"pooled.png\" pooled_img;\n Printf.printf \"Saved: pooled.png (%dx%d -> %dx%d)\\n\" h w pool_h pool_w;\n\n Printf.printf \"\\nAll images saved to the current directory.\\n\"\n"
},
{
"path": "packages/nx/examples/10-machine-learning/README.md",
"content": "# Machine Learning\n\nFour classic ML algorithms built from Nx primitives: SVD, broadcasting,\nreductions, and scalar loops.\n\n| File | Algorithm | Key Nx operations |\n| ---------- | --------- | --------------------------------------------------------- |\n| `pca.ml` | PCA | `svd`, `mean`, `matmul`, `cumsum` |\n| `kmeans.ml`| K-Means | broadcasting, `argmin`, `categorical`, `sq_distances` |\n| `dbscan.ml`| DBSCAN | pairwise distances, `less_equal_s`, boolean `item` in BFS |\n| `tsne.ml` | t-SNE | `exp`, `log`, Student-t kernel, momentum gradient descent |\n\n## Running\n\n```bash\ndune exec nx/examples/10-machine-learning/pca.exe\ndune exec nx/examples/10-machine-learning/kmeans.exe\ndune exec nx/examples/10-machine-learning/dbscan.exe\ndune exec nx/examples/10-machine-learning/tsne.exe\n```\n"
},
{
"path": "packages/nx/examples/10-machine-learning/dbscan.ml",
"content": "(** DBSCAN density-based clustering.\n\n Generate two dense clusters with scattered noise, find clusters using\n neighbourhood density, and report cluster sizes and noise count. *)\n\nopen Nx\n\nlet () =\n let eps = 1.5 in\n let min_samples = 5 in\n\n (* Two tight blobs plus uniform noise *)\n let c1 = add_s (mul_s (randn Float64 [| 80; 2 |]) 0.6) 3.0 in\n let c2 = sub_s (mul_s (randn Float64 [| 80; 2 |]) 0.6) 3.0 in\n let noise = sub_s (mul_s (rand Float64 [| 20; 2 |]) 14.0) 7.0 in\n let data = concatenate ~axis:0 [ c1; c2; noise ] in\n let n = (shape data).(0) in\n Printf.printf \"Data: %d points (eps=%.1f, min_samples=%d)\\n\\n\" n eps\n min_samples;\n\n (* Pairwise Euclidean distance matrix [n, n] *)\n let diff = sub (expand_dims [ 1 ] data) (expand_dims [ 0 ] data) in\n let dist = sqrt (sum ~axes:[ 2 ] (square diff)) in\n\n (* Neighbour adjacency and core-point mask *)\n let neighbours = less_equal_s dist eps in\n let counts = sum ~axes:[ 1 ] (cast Float64 neighbours) in\n let core = greater_equal_s counts (Float.of_int min_samples) in\n\n (* BFS cluster expansion *)\n let labels = Array.make n (-1) in\n let cluster_id = ref 0 in\n for i = 0 to n - 1 do\n if labels.(i) = -1 && item [ i ] core then begin\n let c = !cluster_id in\n incr cluster_id;\n labels.(i) <- c;\n let q = Queue.create () in\n Queue.push i q;\n while not (Queue.is_empty q) do\n let p = Queue.pop q in\n for j = 0 to n - 1 do\n if labels.(j) = -1 && item [ p; j ] neighbours then begin\n labels.(j) <- c;\n if item [ j ] core then Queue.push j q\n end\n done\n done\n end\n done;\n\n let n_clusters = !cluster_id in\n let n_noise =\n Array.fold_left (fun acc l -> if l = -1 then acc + 1 else acc) 0 labels\n in\n Printf.printf \"Clusters found: %d\\n\" n_clusters;\n Printf.printf \"Noise points: %d\\n\\n\" n_noise;\n for c = 0 to n_clusters - 1 do\n let count =\n Array.fold_left (fun acc l -> if l = c then acc + 1 else acc) 0 labels\n in\n Printf.printf \" Cluster %d: %d points\\n\" c count\n done\n"
},
{
"path": "packages/nx/examples/10-machine-learning/dune",
"content": "(executables\n (names pca kmeans dbscan tsne)\n (libraries nx))\n"
},
{
"path": "packages/nx/examples/10-machine-learning/kmeans.ml",
"content": "(** K-means clustering with kmeans++ initialisation.\n\n Generate synthetic blobs, cluster them with Lloyd's algorithm, and report\n centroid positions and inertia. *)\n\nopen Nx\n\n(* Pairwise squared L2 distances: [n, d] x [k, d] -> [n, k] *)\nlet sq_distances a b =\n sum ~axes:[ 2 ] (square (sub (expand_dims [ 1 ] a) (expand_dims [ 0 ] b)))\n\n(* Isotropic Gaussian blobs around given centres. *)\nlet make_blobs ~samples_per_cluster centers =\n let d = (shape centers).(1) in\n let blobs =\n List.init\n (shape centers).(0)\n (fun c ->\n add (randn Float64 [| samples_per_cluster; d |]) (get [ c ] centers))\n in\n shuffle (concatenate ~axis:0 blobs)\n\n(* Kmeans++ initialisation: pick k centres from data. *)\nlet kmeanspp data k =\n let n = (shape data).(0) in\n let d = (shape data).(1) in\n let centroids = zeros Float64 [| k; d |] in\n let idx = Int32.to_int (item [] (randint Int32 ~high:n [||] 0)) in\n set [ 0 ] centroids (get [ idx ] data);\n for c = 1 to k - 1 do\n let current = slice [ R (0, c); A ] centroids in\n let min_d = min ~axes:[ 1 ] (sq_distances data current) in\n let chosen =\n Int32.to_int (item [] (categorical (log (clamp ~min:1e-30 min_d))))\n in\n set [ c ] centroids (get [ chosen ] data)\n done;\n centroids\n\nlet () =\n let true_centers =\n create Float64 [| 3; 2 |] [| 0.0; 0.0; 7.0; 7.0; -5.0; 10.0 |]\n in\n let data = make_blobs ~samples_per_cluster:100 true_centers in\n let n = (shape data).(0) in\n let d = (shape data).(1) in\n let k = 3 in\n Printf.printf \"Data: %d points, %d features, %d clusters\\n\\n\" n d k;\n\n let centroids = kmeanspp data k in\n let labels = ref (zeros Int32 [| n |]) in\n let max_iter = 100 in\n let tol = 1e-6 in\n let converged = ref false in\n let iter = ref 0 in\n while !iter < max_iter && not !converged do\n labels := argmin ~axis:1 (sq_distances data centroids);\n\n let old = copy centroids in\n for c = 0 to k - 1 do\n let mask = cast Float64 (equal !labels (scalar Int32 (Int32.of_int c))) in\n let count = item [] (sum mask) in\n if count > 0.0 then begin\n let total = sum ~axes:[ 0 ] (mul data (expand_dims [ 1 ] mask)) in\n set [ c ] centroids (div_s total count)\n end\n done;\n\n let shift = item [] (max (abs (sub centroids old))) in\n converged := shift < tol;\n incr iter\n done;\n\n Printf.printf \"Converged after %d iterations\\n\\n\" !iter;\n Printf.printf \"Centroids:\\n%s\\n\" (data_to_string centroids);\n\n for c = 0 to k - 1 do\n let count =\n item []\n (sum (cast Float64 (equal !labels (scalar Int32 (Int32.of_int c)))))\n in\n Printf.printf \" Cluster %d: %.0f points\\n\" c count\n done;\n\n let inertia = item [] (sum (min ~axes:[ 1 ] (sq_distances data centroids))) in\n Printf.printf \"\\nInertia: %.2f\\n\" inertia\n"
},
{
"path": "packages/nx/examples/10-machine-learning/pca.ml",
"content": "(** Principal component analysis via SVD.\n\n Generate synthetic data with known structure, project to lower dimensions,\n and verify the explained variance captures the signal. *)\n\nopen Nx\nopen Nx.Infix\n\nlet () =\n (* 200 points in 5D: most variance along axes 0 and 1 *)\n let n = 200 in\n let scale = create Float64 [| 1; 5 |] [| 10.0; 5.0; 1.0; 1.0; 1.0 |] in\n let data = randn Float64 [| n; 5 |] * scale in\n Printf.printf \"Data shape: [%d; %d]\\n\\n\" (shape data).(0) (shape data).(1);\n\n (* Center *)\n let mu = mean ~axes:[ 0 ] ~keepdims:true data in\n let centered = data - mu in\n\n (* Economy SVD: centered = U diag(S) Vt *)\n let _u, s, vt = svd ~full_matrices:false centered in\n\n (* Explained variance ratio: s_i^2 / sum(s^2) *)\n let s2 = square s in\n let ratios = s2 /$ item [] (sum s2) in\n Printf.printf \"Singular values: %s\\n\" (data_to_string s);\n Printf.printf \"Explained variance ratio: %s\\n\" (data_to_string ratios);\n Printf.printf \"Cumulative: %s\\n\\n\"\n (data_to_string (cumsum ratios));\n\n (* Project to 2 components *)\n let n_components = 2 in\n let components = slice [ R (0, n_components); A ] vt in\n let projected = matmul centered (matrix_transpose components) in\n Printf.printf \"Projected shape: [%d; %d]\\n\"\n (shape projected).(0)\n (shape projected).(1);\n\n (* Reconstruct and measure error *)\n let reconstructed = matmul projected components + mu in\n let rmse = sqrt (mean (square (data - reconstructed))) in\n Printf.printf \"Reconstruction RMSE (2 of 5 components): %.4f\\n\" (item [] rmse)\n"
},
{
"path": "packages/nx/examples/10-machine-learning/tsne.ml",
"content": "(** t-SNE dimensionality reduction.\n\n Embed three 10-dimensional clusters into 2D using the exact t-SNE algorithm.\n Reports KL divergence and per-cluster spread. *)\n\nopen Nx\n\n(* Pairwise squared distances: [n, d] -> [n, n] *)\nlet pairwise_sq data =\n let diff = sub (expand_dims [ 1 ] data) (expand_dims [ 0 ] data) in\n sum ~axes:[ 2 ] (square diff)\n\n(* Off-diagonal mask: 1 everywhere except the diagonal. *)\nlet off_diag n =\n sub (full Float64 [| n; n |] 1.0) (cast Float64 (eye Float64 n))\n\n(* Compute symmetric P matrix via binary search for each row's bandwidth. *)\nlet compute_p dist_sq ~perplexity =\n let n = (shape dist_sq).(0) in\n let target = Float.log perplexity in\n let p = zeros Float64 [| n; n |] in\n for i = 0 to n - 1 do\n let di = get [ i ] dist_sq in\n let lo = ref 1e-10 in\n let hi = ref 1e4 in\n let row = ref (zeros Float64 [| n |]) in\n for _ = 0 to 50 do\n let sigma = (!lo +. !hi) /. 2.0 in\n let beta = 1.0 /. (2.0 *. sigma *. sigma) in\n let pi = exp (mul_s di (-.beta)) in\n set_item [ i ] 0.0 pi;\n let s = item [] (sum pi) in\n let pi = div_s pi (Float.max s 1e-30) in\n let h = -.item [] (sum (mul pi (log (clamp ~min:1e-30 pi)))) in\n row := pi;\n if h > target then hi := sigma else lo := sigma\n done;\n set [ i ] p !row\n done;\n (* Symmetrise: P = (P + P^T) / (2n) *)\n let p = div_s (add p (matrix_transpose p)) (2.0 *. Float.of_int n) in\n clamp ~min:1e-12 p\n\nlet () =\n let n_per = 50 in\n let dim = 10 in\n let perplexity = 20.0 in\n let max_iter = 500 in\n let lr = 100.0 in\n\n (* Three well-separated clusters in 10D *)\n let c0 = randn Float64 [| n_per; dim |] in\n let c1 = add_s (randn Float64 [| n_per; dim |]) 8.0 in\n let c2 = sub_s (randn Float64 [| n_per; dim |]) 8.0 in\n let data = concatenate ~axis:0 [ c0; c1; c2 ] in\n let n = (shape data).(0) in\n Printf.printf \"Data: %d points in %dD, perplexity=%.0f\\n\\n\" n dim perplexity;\n\n let dist_sq = pairwise_sq data in\n let p = compute_p dist_sq ~perplexity in\n\n let y = ref (mul_s (randn Float64 [| n; 2 |]) 1e-4) in\n let vel = ref (zeros Float64 [| n; 2 |]) in\n let mask = off_diag n in\n\n for iter = 1 to max_iter do\n let y_diff = sub (expand_dims [ 1 ] !y) (expand_dims [ 0 ] !y) in\n let y_dsq = sum ~axes:[ 2 ] (square y_diff) in\n let inv_d = mul (div (scalar Float64 1.0) (add_s y_dsq 1.0)) mask in\n let q_sum = Float.max (item [] (sum inv_d)) 1e-30 in\n let q = clamp ~min:1e-12 (div_s inv_d q_sum) in\n\n let p_eff = if iter <= 100 then mul_s p 4.0 else p in\n\n (* Gradient: 4 sum_j (p_ij - q_ij)(y_i - y_j)(1+||y_i-y_j||^2)^{-1} *)\n let mult = mul (sub p_eff q) inv_d in\n let grad =\n mul_s (sum ~axes:[ 1 ] (mul (expand_dims [ 2 ] mult) y_diff)) 4.0\n in\n\n let momentum = if iter <= 100 then 0.5 else 0.8 in\n vel := sub (mul_s !vel momentum) (mul_s grad lr);\n y := add !y !vel;\n\n if iter = 1 || iter mod 100 = 0 then begin\n let kl = item [] (sum (mul p (log (div p q)))) in\n Printf.printf \" iter %4d KL = %.4f\\n\" iter kl\n end\n done;\n\n Printf.printf \"\\nPer-cluster spread (mean std of embedded coordinates):\\n\";\n for c = 0 to 2 do\n let lo = c * n_per in\n let cluster = slice [ R (lo, lo + n_per); A ] !y in\n let sx = item [] (mean (std ~axes:[ 0 ] cluster)) in\n Printf.printf \" Cluster %d: %.4f\\n\" c sx\n done\n"
},
{
"path": "packages/nx/examples/README.md",
"content": "# Nx Examples\n\nTen standalone examples that teach Nx from the ground up. Each builds on the\nprevious one, progressing from array creation to machine learning.\n\n## Examples\n\n| # | Example | What You'll Learn |\n| --- | ------------------------------------------------------------ | -------------------------------------------------------------------- |\n| 01 | [Creating Arrays](01-creating-arrays/) | `zeros`, `ones`, `arange`, `linspace`, `init`, `meshgrid`, dtypes |\n| 02 | [Infix and Arithmetic](02-infix-and-arithmetic/) | `Nx.Infix` operators (`+`, `*$`, `/`), `abs`, `sqrt`, `exp`, `clamp` |\n| 03 | [Indexing and Slicing](03-indexing-and-slicing/) | `I`, `R`, `Rs`, `A`, `.${[...]}`, `compress`, `where`, `take` |\n| 04 | [Reshaping and Broadcasting](04-reshaping-and-broadcasting/) | `reshape`, `flatten`, `transpose`, `vstack`, broadcasting rules |\n| 05 | [Reductions and Statistics](05-reductions-and-statistics/) | `mean`, `std`, `argmax`, `cumsum`, `all`, `any`, axis parameter |\n| 06 | [Random Numbers](06-random-numbers/) | `Rng.run`, `Rng.uniform`, `Rng.normal`, `Rng.shuffle`, Monte Carlo |\n| 07 | [Linear Algebra](07-linear-algebra/) | `@@`, `/@`, `inv`, `det`, `lstsq`, `eigh`, `svd` |\n| 08 | [Signal Processing](08-signal-processing/) | `rfft`, `irfft`, `rfftfreq`, frequency filtering |\n| 09 | [Image Processing](09-image-processing/) | `correlate2d`, `max_pool2d`, Sobel edges, `Nx_io.save_image` |\n| 10 | [Machine Learning](10-machine-learning/) | PCA, K-Means, DBSCAN, t-SNE from Nx primitives |\n\n## Running\n\nFrom the repository root:\n\n```bash\ndune exec nx/examples/01-creating-arrays/main.exe\ndune exec nx/examples/02-infix-and-arithmetic/main.exe\n# ... and so on through 10\n```\n\n## Dependencies\n\n- Examples 01-08 and 10 use only `nx`\n- Example 09 adds `nx.io` (image I/O)\n"
},
{
"path": "packages/nx/lib/.ocamlformat-ignore",
"content": "prelude.ml"
},
{
"path": "packages/nx/lib/backend/dune",
"content": "(library\n (name nx_backend)\n (public_name nx.backend)\n (virtual_modules nx_backend)\n (default_implementation nx_c)\n (libraries nx_core nx_buffer))\n"
},
{
"path": "packages/nx/lib/backend/nx_backend.mli",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\ninclude Nx_core.Backend_intf.S\n\n(* TODO: [create_context : unit -> context] won't work for backends that need\n parameters (e.g. GPU device index, memory limits). We'll likely need\n [create_context : ?config:config -> unit -> context] or similar, with a\n default config for each backend. *)\nval create_context : unit -> context\n"
},
{
"path": "packages/nx/lib/backend_c/config/discover.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nmodule C = Configurator.V1\n\nlet test_blas =\n {|\n#include \n\nint main() {\n float a[6] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};\n float b[6] = {7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f};\n float c[9] = {0.0f};\n cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, 3, 3, 2, 1.0f, a, 2,\n b, 3, 0.0f, c, 3);\n return (int)c[0];\n}\n|}\n\nlet test_lapacke =\n {|\n#include \n\nint main() {\n double a[4] = {4.0, 1.0, 1.0, 3.0};\n lapack_int info = LAPACKE_dpotrf(LAPACK_ROW_MAJOR, 'L', 2, a, 2);\n return (int)info;\n}\n|}\n\nlet openblas_default system : C.Pkg_config.package_conf =\n match system with\n | \"mingw\" | \"mingw64\" ->\n (* On cygwin, discover.exe is a mingw binary whose Sys.file_exists cannot\n resolve cygwin paths. The C compiler is a cygwin binary and handles\n them fine, so pass the known sysroot paths directly. *)\n let triplet =\n if system = \"mingw64\" then \"x86_64-w64-mingw32\" else \"i686-w64-mingw32\"\n in\n let prefix = \"/usr/\" ^ triplet ^ \"/sys-root/mingw\" in\n {\n C.Pkg_config.cflags = [ \"-I\" ^ prefix ^ \"/include\" ];\n libs = [ \"-L\" ^ prefix ^ \"/lib\"; \"-lcblas\" ];\n }\n | _ ->\n let search =\n [ \"/usr/local/opt/openblas/lib\"; \"/opt/OpenBLAS/lib/\"; \"/usr/lib\" ]\n |> List.filter Sys.file_exists\n in\n let libs = List.map (fun path -> \"-L\" ^ path) search @ [ \"-lopenblas\" ] in\n let include_dirs =\n [ \"/usr/include/openblas\" ] |> List.filter Sys.file_exists\n in\n let cflags = List.map (fun path -> \"-I\" ^ path) include_dirs in\n { C.Pkg_config.cflags; libs }\n\nlet default_ldlibs = [ \"-lm\" ]\n\nlet string_contains ~needle haystack =\n let h_len = String.length haystack in\n let n_len = String.length needle in\n let rec aux i =\n if i + n_len > h_len then false\n else if String.sub haystack i n_len = needle then true\n else aux (i + 1)\n in\n if n_len = 0 then true else aux 0\n\nlet list_find_map f lst =\n let rec aux = function\n | [] -> None\n | x :: xs -> ( match f x with None -> aux xs | some -> some)\n in\n aux lst\n\nlet libomp_paths c =\n let env = Sys.getenv_opt \"LIBOMP_PREFIX\" in\n let brew_prefix =\n if C.Process.run_ok c \"brew\" [ \"--prefix\"; \"libomp\" ] then\n Some\n (C.Process.run_capture_exn c \"brew\" [ \"--prefix\"; \"libomp\" ]\n |> String.trim)\n else None\n in\n let candidates =\n List.filter_map\n (fun x -> x)\n [\n env;\n brew_prefix;\n Some \"/opt/homebrew/opt/libomp\";\n Some \"/usr/local/opt/libomp\";\n ]\n in\n list_find_map\n (fun prefix ->\n let include_dir = Filename.concat prefix \"include\" in\n let lib_dir = Filename.concat prefix \"lib\" in\n let header = Filename.concat include_dir \"omp.h\" in\n if Sys.file_exists header then\n Some ([ \"-I\" ^ include_dir ], [ \"-L\" ^ lib_dir ])\n else None)\n candidates\n\nlet compiler_is_clang c =\n let compiler =\n Sys.getenv_opt \"CC\"\n |> Option.value\n ~default:\n (match C.ocaml_config_var c \"c_compiler\" with\n | Some cc when String.trim cc <> \"\" -> cc\n | _ -> \"cc\")\n in\n if C.Process.run_ok c compiler [ \"--version\" ] then\n let version =\n C.Process.run_capture_exn c compiler [ \"--version\" ]\n |> String.lowercase_ascii\n in\n string_contains ~needle:\"clang\" version\n else false\n\nlet detect_openmp c system base_flags =\n let test c_flags link_flags =\n C.c_test c ~c_flags ~link_flags\n \"#include \\nint main(){return omp_get_num_threads();}\"\n in\n match system with\n | \"macosx\" ->\n if compiler_is_clang c then\n let include_flags, lib_dir_flags =\n match libomp_paths c with Some paths -> paths | None -> ([], [])\n in\n let openmp_flags = include_flags @ [ \"-Xpreprocessor\"; \"-fopenmp\" ] in\n let openmp_libs = lib_dir_flags @ [ \"-lomp\" ] in\n if test openmp_flags openmp_libs then\n (base_flags @ openmp_flags, openmp_libs)\n else (base_flags, [])\n else if test [ \"-fopenmp\" ] [ \"-fopenmp\" ] then\n (base_flags @ [ \"-fopenmp\" ], [ \"-fopenmp\" ])\n else (base_flags, [])\n | \"linux\" | \"linux_elf\" | \"mingw\" | \"mingw64\" | \"cygwin\" ->\n if test [ \"-fopenmp\" ] [ \"-fopenmp\" ] then\n (base_flags @ [ \"-fopenmp\" ], [ \"-fopenmp\" ])\n else (base_flags, [])\n | _ ->\n if test [ \"-fopenmp\" ] [ \"-fopenmp\" ] then\n (base_flags @ [ \"-fopenmp\" ], [ \"-fopenmp\" ])\n else (base_flags, [])\n\nlet pkg_query c package =\n match C.Pkg_config.get c with\n | None -> None\n | Some pc -> C.Pkg_config.query pc ~package\n\nlet ensure_lapacke c c_flags libs pkg_query_fn =\n if C.c_test c test_lapacke ~c_flags ~link_flags:libs then (libs, c_flags)\n else\n let lapacke_conf =\n match pkg_query_fn \"llapacke\" with\n | Some conf -> conf\n | None -> (\n match pkg_query_fn \"lapacke\" with\n | Some conf -> conf\n | None -> { C.Pkg_config.cflags = []; libs = [ \"-llapacke\" ] })\n in\n let libs = lapacke_conf.libs @ libs in\n let c_flags = lapacke_conf.cflags @ c_flags in\n if C.c_test c test_lapacke ~c_flags ~link_flags:libs then (libs, c_flags)\n else (\n Printf.printf\n {|\nUnable to link against LAPACKE even after adding (%s) to the link flags.\nVerify that a LAPACKE implementation is installed and visible to the\nbuild system (consider installing lapacke or setting PKG_CONFIG_PATH).\n|}\n (String.concat \" \" lapacke_conf.libs);\n failwith \"Unable to link against lapacke.\")\n\nlet () =\n C.main ~name:\"nx_c\" (fun c ->\n let system = C.ocaml_config_var_exn c \"system\" in\n let architecture = C.ocaml_config_var_exn c \"architecture\" in\n let word_size = C.ocaml_config_var_exn c \"word_size\" in\n\n let base_flags =\n let opt_flags =\n match architecture with\n | \"amd64\" | \"x86_64\" -> [ \"-O3\"; \"-march=native\"; \"-fPIC\" ]\n | \"arm64\" | \"aarch64\" -> [ \"-O3\"; \"-mcpu=native\"; \"-fPIC\" ]\n | \"power\" | \"ppc\" | \"ppc64\" | \"ppc64le\" ->\n [ \"-O3\"; \"-mcpu=native\"; \"-fPIC\" ]\n | \"riscv32\" -> [ \"-O3\"; \"-march=rv32gc\"; \"-fPIC\" ]\n | \"riscv64\" -> [ \"-O3\"; \"-march=rv64gc\"; \"-fPIC\" ]\n | \"riscv\" ->\n if word_size = \"64\" then [ \"-O3\"; \"-march=rv64gc\"; \"-fPIC\" ]\n else [ \"-O3\"; \"-march=rv32gc\"; \"-fPIC\" ]\n | \"s390x\" -> [ \"-O3\"; \"-march=native\"; \"-fPIC\" ]\n | _ -> [ \"-O3\"; \"-fPIC\" ]\n in\n (* Suppress vectorization failure warnings from clang *)\n if compiler_is_clang c then opt_flags @ [ \"-Wno-pass-failed\" ]\n else opt_flags\n in\n\n let opt_flags =\n match system with\n | \"macosx\" -> List.filter (fun flag -> flag <> \"-fPIC\") base_flags\n | _ -> base_flags\n in\n\n let opt_flags, openmp_libs = detect_openmp c system opt_flags in\n\n let openblas_conf =\n match pkg_query c \"openblas\" with\n | Some conf -> conf\n | None -> (\n match pkg_query c \"cblas\" with\n | Some conf -> conf\n | None -> openblas_default system)\n in\n\n let filter_openmp_flags flags =\n let rec loop acc = function\n | [] -> List.rev acc\n | \"-Xpreprocessor\" :: \"-fopenmp\" :: rest -> loop acc rest\n | \"-fopenmp\" :: rest -> loop acc rest\n | flag :: rest -> loop (flag :: acc) rest\n in\n loop [] flags\n in\n let openblas_cflags = filter_openmp_flags openblas_conf.cflags in\n let openblas_libs = filter_openmp_flags openblas_conf.libs in\n let c_flags = opt_flags @ openblas_cflags in\n let libs =\n (if system = \"macosx\" then [ \"-framework\"; \"Accelerate\" ] else [])\n @ openblas_libs @ openmp_libs @ default_ldlibs\n in\n\n if not (C.c_test c test_blas ~c_flags ~link_flags:libs) then (\n Printf.printf\n {|\nUnable to link against OpenBLAS: the current values for cflags and libs\nare respectively (%s) and (%s).\nCheck that OpenBLAS is installed and, if necessary, extend PKG_CONFIG_PATH\nwith the directory containing openblas.pc.\n|}\n (String.concat \" \" openblas_cflags)\n (String.concat \" \" openblas_libs);\n failwith \"Unable to link against openblas.\");\n\n let libs, c_flags = ensure_lapacke c c_flags libs (pkg_query c) in\n\n C.Flags.write_sexp \"c_flags.sexp\" c_flags;\n C.Flags.write_sexp \"c_library_flags.sexp\" libs)\n"
},
{
"path": "packages/nx/lib/backend_c/config/dune",
"content": "(executable\n (name discover)\n (libraries dune-configurator))\n"
},
{
"path": "packages/nx/lib/backend_c/dune",
"content": "(library\n (name nx_c)\n (public_name nx.c)\n (implements nx.backend)\n (libraries nx_buffer nx_core nx.pocketfft)\n (foreign_stubs\n (language c)\n (names\n nx_c_binary\n nx_c_unary\n nx_c_reduce\n nx_c_sort\n nx_c_scan\n nx_c_ternary\n nx_c_cast\n nx_c_memory\n nx_c_index\n nx_c_random\n nx_c_shape\n nx_c_window\n nx_c_matmul\n nx_c_cholesky\n nx_c_qr\n nx_c_eig\n nx_c_solve\n nx_c_svd)\n (flags\n :standard\n (:include c_flags.sexp)))\n (c_library_flags\n :standard\n (:include c_library_flags.sexp)))\n\n(rule\n (targets c_library_flags.sexp c_flags.sexp)\n (action\n (run config/discover.exe)))\n"
},
{
"path": "packages/nx/lib/backend_c/nx_backend.ml",
"content": "(*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*)\n\nopen Nx_core\n\nlet err op fmt = Printf.ksprintf (fun msg -> invalid_arg (op ^ \": \" ^ msg)) fmt\n\ntype ('a, 'b) buffer = ('a, 'b) Nx_buffer.t\ntype context = unit\n\nlet create_context () = ()\n\ntype ('a, 'b) t = {\n context : context;\n dtype : ('a, 'b) Dtype.t;\n buffer : ('a, 'b) buffer;\n view : View.t;\n}\n\n(* We define an FFI tensor type for easy access to the view fields in C.\n\n XXX: probably more efficient to inline those in our [t] type and have the\n view function create a view when called. *)\ntype ('a, 'b) ffi_tensor = {\n data : ('a, 'b) buffer;\n shape : int array;\n strides : int array;\n offset : int;\n}\n[@@warning \"-69\"]\n\n(* ───── External FFI Declarations ───── *)\n\nexternal caml_add :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_add\"\n\nexternal caml_mul :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_mul\"\n\nexternal caml_idiv :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_idiv\"\n\nexternal caml_fdiv :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_fdiv\"\n\nexternal caml_max :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_max\"\n\nexternal caml_min :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_min\"\n\nexternal caml_sub :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_sub\"\n\nexternal caml_mod :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_mod\"\n\nexternal caml_pow :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_pow\"\n\nexternal caml_cmpeq :\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n (bool, Dtype.bool_elt) ffi_tensor ->\n unit = \"caml_nx_cmpeq\"\n\nexternal caml_cmpne :\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n (bool, Dtype.bool_elt) ffi_tensor ->\n unit = \"caml_nx_cmpne\"\n\nexternal caml_cmplt :\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n (bool, Dtype.bool_elt) ffi_tensor ->\n unit = \"caml_nx_cmplt\"\n\nexternal caml_cmple :\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n (bool, Dtype.bool_elt) ffi_tensor ->\n unit = \"caml_nx_cmple\"\n\nexternal caml_xor :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_xor\"\n\nexternal caml_or :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_or\"\n\nexternal caml_and :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_and\"\n\nexternal caml_atan2 :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_atan2\"\n\n(* ───── Unary Operation FFI Declarations ───── *)\n\nexternal caml_neg : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_neg\"\n\nexternal caml_sin : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_sin\"\n\nexternal caml_cos : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_cos\"\n\nexternal caml_sqrt : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_sqrt\"\n\nexternal caml_abs : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_abs\"\n\nexternal caml_log : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_log\"\n\nexternal caml_exp : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_exp\"\n\nexternal caml_recip : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_recip\"\n\nexternal caml_sign : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_sign\"\n\nexternal caml_tan : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_tan\"\n\nexternal caml_asin : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_asin\"\n\nexternal caml_acos : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_acos\"\n\nexternal caml_atan : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_atan\"\n\nexternal caml_sinh : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_sinh\"\n\nexternal caml_cosh : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_cosh\"\n\nexternal caml_tanh : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_tanh\"\n\nexternal caml_trunc : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_trunc\"\n\nexternal caml_ceil : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_ceil\"\n\nexternal caml_floor : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_floor\"\n\nexternal caml_round : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_round\"\n\nexternal caml_erf : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_erf\"\n\n(* ───── Ternary Operation FFI Declarations ───── *)\n\nexternal caml_where :\n (bool, Dtype.bool_elt) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n unit = \"caml_nx_where\"\n\n(* ───── Reduction Operation FFI Declarations ───── *)\n\nexternal caml_reduce_sum :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> int array -> bool -> unit\n = \"caml_nx_reduce_sum\"\n\nexternal caml_reduce_max :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> int array -> bool -> unit\n = \"caml_nx_reduce_max\"\n\nexternal caml_reduce_prod :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> int array -> bool -> unit\n = \"caml_nx_reduce_prod\"\n\nexternal caml_reduce_min :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> int array -> bool -> unit\n = \"caml_nx_reduce_min\"\n\nexternal caml_associative_scan :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> int -> int -> unit\n = \"caml_nx_associative_scan\"\n\nexternal caml_argmax :\n ('a, 'b) ffi_tensor ->\n (int32, Dtype.int32_elt) ffi_tensor ->\n int ->\n bool ->\n unit = \"caml_nx_argmax\"\n\nexternal caml_argmin :\n ('a, 'b) ffi_tensor ->\n (int32, Dtype.int32_elt) ffi_tensor ->\n int ->\n bool ->\n unit = \"caml_nx_argmin\"\n\nexternal caml_sort :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> int -> bool -> unit\n = \"caml_nx_sort\"\n\nexternal caml_argsort :\n ('a, 'b) ffi_tensor ->\n (int32, Dtype.int32_elt) ffi_tensor ->\n int ->\n bool ->\n unit = \"caml_nx_argsort\"\n\n(* Cast operation FFI declaration *)\nexternal caml_cast : ('a, 'b) ffi_tensor -> ('c, 'd) ffi_tensor -> unit\n = \"caml_nx_cast\"\n\n(* ───── Memory Operation FFI Declarations ───── *)\n\nexternal caml_copy : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor = \"caml_nx_copy\"\n\nexternal caml_contiguous : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor\n = \"caml_nx_contiguous\"\n\nexternal caml_assign : ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_assign\"\n\n(* ───── Index Operation FFI Declarations ───── *)\n\nexternal caml_gather :\n ('a, 'b) ffi_tensor ->\n (int32, Dtype.int32_elt) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n int ->\n unit = \"caml_nx_op_gather\"\n\nexternal caml_scatter :\n ('a, 'b) ffi_tensor ->\n (int32, Dtype.int32_elt) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n int ->\n ('a, 'b) ffi_tensor ->\n int ->\n bool ->\n unit = \"caml_nx_op_scatter_bc\" \"caml_nx_op_scatter\"\n\n(* ───── Linear Algebra Operation FFI Declarations ───── *)\n\nexternal caml_cholesky :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> bool -> unit\n = \"caml_nx_op_cholesky\"\n\nexternal caml_matmul :\n ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_matmul\"\n\nexternal caml_triangular_solve :\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n bool ->\n bool ->\n bool ->\n unit = \"caml_nx_op_triangular_solve_bc\" \"caml_nx_op_triangular_solve\"\n\nexternal caml_qr :\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n bool ->\n unit = \"caml_nx_op_qr\"\n\nexternal caml_eig :\n ('a, 'b) ffi_tensor ->\n ('c, 'd) ffi_tensor ->\n ('e, 'f) ffi_tensor ->\n bool ->\n bool ->\n unit = \"caml_nx_op_eig\"\n\nexternal caml_svd :\n ('a, 'b) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n ('c, 'd) ffi_tensor ->\n ('a, 'b) ffi_tensor ->\n bool ->\n unit = \"caml_nx_op_svd\"\n\n(* ───── Shape Operation FFI Declarations ───── *)\n\nexternal caml_cat :\n ('a, 'b) ffi_tensor list -> int -> ('a, 'b) ffi_tensor -> unit = \"caml_nx_cat\"\n\nexternal caml_pad :\n ('a, 'b) ffi_tensor -> int array -> 'a -> ('a, 'b) ffi_tensor -> unit\n = \"caml_nx_pad\"\n\n(* ───── Window Operation FFI Declarations ───── *)\n\nexternal caml_unfold :\n ('a, 'b) ffi_tensor ->\n int array ->\n int array ->\n int array ->\n int array ->\n ('a, 'b) ffi_tensor ->\n unit = \"caml_nx_op_unfold_bc\" \"caml_nx_op_unfold\"\n\nexternal caml_fold :\n ('a, 'b) ffi_tensor ->\n int array ->\n int array ->\n int array ->\n int array ->\n int array ->\n ('a, 'b) ffi_tensor ->\n unit = \"caml_nx_op_fold_bc\" \"caml_nx_op_fold\"\n\n(* ───── Random Operation FFI Declarations ───── *)\n\nexternal caml_threefry :\n (int32, Dtype.int32_elt) ffi_tensor ->\n (int32, Dtype.int32_elt) ffi_tensor ->\n (int32, Dtype.int32_elt) ffi_tensor ->\n unit = \"caml_nx_threefry\"\n\n(* ───── Helper Functions ───── *)\n\nlet view t = t.view\nlet dtype t = t.dtype\nlet to_host t = t.buffer\nlet context t = t.context\nlet shape t = View.shape t.view\n\nlet strides t =\n match View.strides_opt t.view with\n | Some s -> s\n | None -> invalid_arg \"strides: cannot get strides for view\"\n\nlet offset t = View.offset t.view\nlet is_contiguous t = View.is_c_contiguous t.view\n\n(* Check if a tensor can be efficiently operated on *)\nlet can_get_strides t = View.can_get_strides t.view\n\n(* Convert tensor to FFI representation if possible *)\nlet to_ffi_tensor t =\n if not (can_get_strides t) then\n invalid_arg \"to_ffi_tensor: tensor has non-materializable view\"\n else\n { data = t.buffer; shape = shape t; strides = strides t; offset = offset t }\n\n(* Create a new tensor with given shape *)\nlet create_tensor ctx dtype shape_arr =\n let size = Array.fold_left ( * ) 1 shape_arr in\n let kind = Dtype.to_buffer_kind dtype in\n let buffer = Nx_buffer.create kind size in\n let view = View.create shape_arr in\n { context = ctx; dtype; buffer; view }\n\nlet buffer ctx dtype shape_arr =\n let kind = Dtype.to_buffer_kind dtype in\n let size = Array.fold_left ( * ) 1 shape_arr in\n let buffer = Nx_buffer.create kind size in\n let view = View.create shape_arr in\n { context = ctx; dtype; buffer; view }\n\nlet full ctx dtype shape_arr value =\n let t = buffer ctx dtype shape_arr in\n Nx_buffer.fill t.buffer value;\n t\n\n(* Materialize a tensor to contiguous layout if needed *)\nlet materialize t =\n (* Check if it has broadcast dimensions (zero strides) *)\n let strides_arr = strides t in\n let has_broadcast = Array.exists (( = ) 0) strides_arr in\n\n if is_contiguous t && offset t = 0 && not has_broadcast then t\n else\n (* Create a contiguous copy *)\n let out_shape = shape t in\n let out = create_tensor t.context t.dtype out_shape in\n let t_ffi = to_ffi_tensor t in\n let out_ffi = to_ffi_tensor out in\n caml_assign t_ffi out_ffi;\n out\n\n(* Ensure tensor is materializable for C operations *)\nlet ensure_materializable t =\n if not (can_get_strides t) then\n (* Broadcast views or complex chains need materialization *)\n materialize t\n else\n (* Check for zero strides (broadcast dimensions) *)\n let strides_arr = strides t in\n if Array.exists (( = ) 0) strides_arr then\n (* Has broadcast dimensions - need to materialize *)\n materialize t\n else t\n\n(* Generic binary operation - allocates output and returns it *)\nlet binary_op op_name ffi_op x y =\n let x_shape = shape x in\n let y_shape = shape y in\n if x_shape <> y_shape then\n err op_name \"shape mismatch: x %s, y %s\" (Shape.to_string x_shape)\n (Shape.to_string y_shape)\n else\n let x' = ensure_materializable x in\n let y' = ensure_materializable y in\n let out = buffer () (dtype x) x_shape in\n let x_ffi = to_ffi_tensor x' in\n let y_ffi = to_ffi_tensor y' in\n let out_ffi = to_ffi_tensor out in\n ffi_op x_ffi y_ffi out_ffi;\n out\n\n(* Comparison operation - allocates bool output and returns it *)\nlet comparison_op op_name ffi_op x y =\n let x_shape = shape x in\n let y_shape = shape y in\n if x_shape <> y_shape then\n err op_name \"shape mismatch: x %s, y %s\" (Shape.to_string x_shape)\n (Shape.to_string y_shape)\n else\n let x' = ensure_materializable x in\n let y' = ensure_materializable y in\n let out = buffer () Dtype.Bool x_shape in\n let x_ffi = to_ffi_tensor x' in\n let y_ffi = to_ffi_tensor y' in\n let out_ffi = to_ffi_tensor out in\n ffi_op x_ffi y_ffi out_ffi;\n out\n\n(* ───── Buffer Allocation ───── *)\n\nlet from_host ctx array =\n let dtype = Dtype.of_buffer_kind (Nx_buffer.kind array) in\n let size = Nx_buffer.length array in\n (* Create a view for the 1D array *)\n let view = View.create [| size |] in\n (* Note: We're sharing the buffer directly, assuming it's contiguous *)\n { context = ctx; dtype; buffer = array; view }\n\n(* Generic unary operation - allocates output and returns it *)\nlet unary_op _op_name ffi_op x =\n let x' = ensure_materializable x in\n let out = buffer () (dtype x) (shape x) in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n ffi_op x_ffi out_ffi;\n out\n\n(* ───── Binary Operations ───── *)\n\nlet add x y = binary_op \"add\" caml_add x y\nlet sub x y = binary_op \"sub\" caml_sub x y\nlet mul x y = binary_op \"mul\" caml_mul x y\nlet max x y = binary_op \"max\" caml_max x y\nlet min x y = binary_op \"min\" caml_min x y\nlet mod_ x y = binary_op \"mod\" caml_mod x y\nlet pow x y = binary_op \"pow\" caml_pow x y\nlet xor x y = binary_op \"xor\" caml_xor x y\nlet or_ x y = binary_op \"or\" caml_or x y\nlet and_ x y = binary_op \"and\" caml_and x y\nlet atan2 y x = binary_op \"atan2\" caml_atan2 y x\n\n(* ───── Comparison Operations ───── *)\n\nlet cmpeq x y = comparison_op \"cmpeq\" caml_cmpeq x y\nlet cmpne x y = comparison_op \"cmpne\" caml_cmpne x y\nlet cmplt x y = comparison_op \"cmplt\" caml_cmplt x y\nlet cmple x y = comparison_op \"cmple\" caml_cmple x y\n\n(* ───── Unary Operations ───── *)\n\nlet neg x = unary_op \"neg\" caml_neg x\nlet log x = unary_op \"log\" caml_log x\nlet exp x = unary_op \"exp\" caml_exp x\nlet sin x = unary_op \"sin\" caml_sin x\nlet cos x = unary_op \"cos\" caml_cos x\nlet sqrt x = unary_op \"sqrt\" caml_sqrt x\nlet abs x = unary_op \"abs\" caml_abs x\nlet recip x = unary_op \"recip\" caml_recip x\nlet sign x = unary_op \"sign\" caml_sign x\nlet tan x = unary_op \"tan\" caml_tan x\nlet asin x = unary_op \"asin\" caml_asin x\nlet acos x = unary_op \"acos\" caml_acos x\nlet atan x = unary_op \"atan\" caml_atan x\nlet sinh x = unary_op \"sinh\" caml_sinh x\nlet cosh x = unary_op \"cosh\" caml_cosh x\nlet tanh x = unary_op \"tanh\" caml_tanh x\nlet trunc x = unary_op \"trunc\" caml_trunc x\nlet ceil x = unary_op \"ceil\" caml_ceil x\nlet floor x = unary_op \"floor\" caml_floor x\nlet round x = unary_op \"round\" caml_round x\nlet erf x = unary_op \"erf\" caml_erf x\n\n(* Ternary Op - allocates output and returns it *)\nlet where cond if_true if_false =\n let cond_shape = shape cond in\n let if_true_shape = shape if_true in\n let if_false_shape = shape if_false in\n\n if cond_shape <> if_true_shape || if_true_shape <> if_false_shape then\n err \"where\" \"shape mismatch: cond %s, if_true %s, if_false %s\"\n (Shape.to_string cond_shape)\n (Shape.to_string if_true_shape)\n (Shape.to_string if_false_shape)\n else\n let cond' = ensure_materializable cond in\n let if_true' = ensure_materializable if_true in\n let if_false' = ensure_materializable if_false in\n let out = buffer () (dtype if_true) if_true_shape in\n let cond_ffi = to_ffi_tensor cond' in\n let if_true_ffi = to_ffi_tensor if_true' in\n let if_false_ffi = to_ffi_tensor if_false' in\n let out_ffi = to_ffi_tensor out in\n caml_where cond_ffi if_true_ffi if_false_ffi out_ffi;\n out\n\n(* Reduction Ops - allocates output and returns it *)\nlet reduce_op _op_name ffi_op ~axes ~keepdims x =\n let input_shape = shape x in\n let ndim = Array.length input_shape in\n\n if ndim = 0 then begin\n let out = buffer () (dtype x) [||] in\n Nx_buffer.set out.buffer 0 (Nx_buffer.get x.buffer 0);\n out\n end\n else\n let normalized_axes =\n Array.map (fun ax -> if ax < 0 then ax + ndim else ax) axes\n in\n let out_shape =\n Shape.reduce_output_shape input_shape normalized_axes keepdims\n in\n let out = buffer () (dtype x) out_shape in\n let x' = ensure_materializable x in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n ffi_op x_ffi out_ffi normalized_axes keepdims;\n out\n\nlet reduce_sum ~axes ~keepdims x =\n reduce_op \"reduce_sum\" caml_reduce_sum ~axes ~keepdims x\n\nlet reduce_max ~axes ~keepdims x =\n reduce_op \"reduce_max\" caml_reduce_max ~axes ~keepdims x\n\nlet reduce_prod ~axes ~keepdims x =\n reduce_op \"reduce_prod\" caml_reduce_prod ~axes ~keepdims x\n\nlet reduce_min ~axes ~keepdims x =\n reduce_op \"reduce_min\" caml_reduce_min ~axes ~keepdims x\n\nlet associative_scan ~axis ~op x =\n let x_shape = shape x in\n let rank = Array.length x_shape in\n if rank = 0 then invalid_arg \"associative_scan: requires rank >= 1\"\n else\n let axis = if axis < 0 then axis + rank else axis in\n if axis < 0 || axis >= rank then\n err \"associative_scan\" \"axis %d out of bounds for rank %d\" axis rank\n else\n let x' = ensure_materializable x in\n let out = buffer () (dtype x) x_shape in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n let op_tag =\n match op with `Sum -> 0 | `Prod -> 1 | `Max -> 2 | `Min -> 3\n in\n caml_associative_scan x_ffi out_ffi axis op_tag;\n out\n\n(* Movement Ops - These are view-only operations *)\nlet expand x shape = { x with view = View.expand x.view shape }\nlet reshape x shape = { x with view = View.reshape x.view shape }\nlet permute x axes = { x with view = View.permute x.view axes }\n\nlet pad x padding fill_value =\n let x' = ensure_materializable x in\n\n (* Calculate output shape *)\n let in_shape = shape x in\n let ndim = Array.length in_shape in\n\n (* Convert pairs to flat array for C interface *)\n let padding_flat =\n Array.init (2 * ndim) (fun i ->\n let dim = i / 2 in\n if i mod 2 = 0 then fst padding.(dim) else snd padding.(dim))\n in\n\n (* Calculate output shape *)\n let out_shape =\n Array.init ndim (fun i ->\n let before, after = padding.(i) in\n in_shape.(i) + before + after)\n in\n\n let out = create_tensor x.context x.dtype out_shape in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n caml_pad x_ffi padding_flat fill_value out_ffi;\n out\n\nlet shrink x bounds = { x with view = View.shrink x.view bounds }\nlet flip x axes = { x with view = View.flip x.view axes }\n\nlet cat tensors ~axis =\n match tensors with\n | [] -> invalid_arg \"cat: empty tensor list\"\n | first :: _ ->\n let tensors' = List.map ensure_materializable tensors in\n\n (* Calculate output shape *)\n let first_shape = shape first in\n let ndim = Array.length first_shape in\n let norm_axis = if axis < 0 then ndim + axis else axis in\n\n (* Sum up dimensions along concatenation axis *)\n let total_axis_size =\n List.fold_left\n (fun acc t ->\n let s = shape t in\n acc + s.(norm_axis))\n 0 tensors\n in\n\n let out_shape =\n Array.mapi\n (fun i dim -> if i = norm_axis then total_axis_size else dim)\n first_shape\n in\n let out = buffer () (dtype first) out_shape in\n let tensors_ffi = List.map to_ffi_tensor tensors' in\n let out_ffi = to_ffi_tensor out in\n caml_cat tensors_ffi norm_axis out_ffi;\n out\n\n(* ───── Other Ops ───── *)\n\nlet cast (type a b c d) ~(dtype : (c, d) Dtype.t) (x : (a, b) t) : (c, d) t =\n let x' = ensure_materializable x in\n let out = buffer () dtype (shape x) in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n caml_cast x_ffi out_ffi;\n out\n\nlet contiguous x =\n (* Check if already contiguous with no offset and no broadcast dimensions *)\n let strides_arr = strides x in\n let has_broadcast = Array.exists (( = ) 0) strides_arr in\n if is_contiguous x && offset x = 0 && not has_broadcast then x\n else\n let x' = ensure_materializable x in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = caml_contiguous x_ffi in\n (* Create tensor from FFI result - it's contiguous so simple view *)\n let view = View.create out_ffi.shape in\n { context = x.context; dtype = x.dtype; buffer = out_ffi.data; view }\n\nlet copy x =\n let x' = ensure_materializable x in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = caml_copy x_ffi in\n (* Create tensor from FFI result - it's contiguous so simple view *)\n let view = View.create out_ffi.shape in\n { context = x.context; dtype = x.dtype; buffer = out_ffi.data; view }\n\nlet assign dst src =\n let src' = ensure_materializable src in\n (* dst doesn't need materialization - we're writing to it *)\n let src_ffi = to_ffi_tensor src' in\n let dst_ffi = to_ffi_tensor dst in\n caml_assign src_ffi dst_ffi\n\nlet threefry key counter =\n let key' = ensure_materializable key in\n let counter' = ensure_materializable counter in\n let out = buffer () Dtype.Int32 (shape counter) in\n let key_ffi = to_ffi_tensor key' in\n let counter_ffi = to_ffi_tensor counter' in\n let out_ffi = to_ffi_tensor out in\n caml_threefry key_ffi counter_ffi out_ffi;\n out\n\n(* ───── Element Access Ops ───── *)\n\nlet gather data indices ~axis =\n (* Ensure inputs are materializable. Preserve broadcasted strides for indices\n to enable C fast paths (e.g., memcpy row gather). *)\n let data' = ensure_materializable data in\n (* Do not materialize indices unless we cannot get strides *)\n let indices' =\n if can_get_strides indices then indices else ensure_materializable indices\n in\n let out = buffer () (dtype data) (shape indices) in\n let data_ffi = to_ffi_tensor data' in\n let indices_ffi = to_ffi_tensor indices' in\n let out_ffi = to_ffi_tensor out in\n caml_gather data_ffi indices_ffi out_ffi axis;\n out\n\nlet scatter ?(mode = `Set) ?(unique_indices = false) data_template ~indices\n ~updates ~axis =\n (* Ensure inputs are materializable *)\n let template' = ensure_materializable data_template in\n let indices' = ensure_materializable indices in\n let updates' = ensure_materializable updates in\n\n (* Create output tensor - for Set mode, start with a copy of template *)\n let out =\n if mode = `Set then copy data_template (* Start with copy of template *)\n else\n create_tensor data_template.context data_template.dtype\n (shape data_template)\n in\n\n (* Convert to FFI tensors *)\n let template_ffi = to_ffi_tensor template' in\n let indices_ffi = to_ffi_tensor indices' in\n let updates_ffi = to_ffi_tensor updates' in\n let out_ffi = to_ffi_tensor out in\n\n (* Convert mode to integer: 0 for Set, 1 for Add *)\n let mode_int = match mode with `Set -> 0 | `Add -> 1 in\n\n (* Call FFI function *)\n caml_scatter template_ffi indices_ffi updates_ffi axis out_ffi mode_int\n unique_indices;\n\n out\n\nlet unfold x ~kernel_size ~stride ~dilation ~padding =\n let x' = ensure_materializable x in\n let in_shape = shape x in\n let k = Array.length kernel_size in\n let leading_ndim = Array.length in_shape - k in\n let leading_shape = Array.sub in_shape 0 leading_ndim in\n let spatial_dims = Array.sub in_shape leading_ndim k in\n\n let padding_flat =\n Array.init\n (Array.length padding * 2)\n (fun i ->\n let dim = i / 2 in\n if i mod 2 = 0 then fst padding.(dim) else snd padding.(dim))\n in\n\n let out_spatial =\n Array.init k (fun i ->\n let pad_before, pad_after = padding.(i) in\n let padded = spatial_dims.(i) + pad_before + pad_after in\n let kernel_extent = (dilation.(i) * (kernel_size.(i) - 1)) + 1 in\n let diff = padded - kernel_extent in\n if diff < 0 then\n invalid_arg \"unfold: kernel size larger than padded input\"\n else (diff / stride.(i)) + 1)\n in\n\n let kernel_prod = Array.fold_left ( * ) 1 kernel_size in\n let spatial_prod = Array.fold_left ( * ) 1 out_spatial in\n let out_shape =\n Array.concat [ leading_shape; [| kernel_prod; spatial_prod |] ]\n in\n\n let out = create_tensor x.context x.dtype out_shape in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n caml_unfold x_ffi kernel_size stride dilation padding_flat out_ffi;\n out\n\nlet fold x ~output_size ~kernel_size ~stride ~dilation ~padding =\n let x' = ensure_materializable x in\n let in_shape = shape x in\n let leading_ndim = Array.length in_shape - 2 in\n let leading_shape = Array.sub in_shape 0 leading_ndim in\n\n let padding_flat =\n Array.init\n (Array.length padding * 2)\n (fun i ->\n let dim = i / 2 in\n if i mod 2 = 0 then fst padding.(dim) else snd padding.(dim))\n in\n\n let _ =\n Array.init (Array.length output_size) (fun i ->\n let pad_before, pad_after = padding.(i) in\n let padded = output_size.(i) + pad_before + pad_after in\n let kernel_extent = (dilation.(i) * (kernel_size.(i) - 1)) + 1 in\n let diff = padded - kernel_extent in\n if diff < 0 then\n invalid_arg \"fold: kernel size larger than padded output\"\n else (diff / stride.(i)) + 1)\n in\n\n let out_shape = Array.concat [ leading_shape; output_size ] in\n\n let out = create_tensor x.context x.dtype out_shape in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n caml_fold x_ffi output_size kernel_size stride dilation padding_flat out_ffi;\n out\n\nlet matmul x y =\n let x' = if is_contiguous x then x else ensure_materializable x in\n let y' = if is_contiguous y then y else ensure_materializable y in\n let x_shape = shape x in\n let y_shape = shape y in\n let x_ndim = Array.length x_shape in\n let y_ndim = Array.length y_shape in\n let m = x_shape.(x_ndim - 2) in\n let n = y_shape.(y_ndim - 1) in\n let max_ndim = Int.max x_ndim y_ndim in\n let batch_nd = max_ndim - 2 in\n let batch_dims =\n Array.init batch_nd (fun i ->\n let a_idx = i - (max_ndim - x_ndim) in\n let b_idx = i - (max_ndim - y_ndim) in\n let sa = if a_idx >= 0 then x_shape.(a_idx) else 1 in\n let sb = if b_idx >= 0 then y_shape.(b_idx) else 1 in\n Int.max sa sb)\n in\n let out_shape = Array.append batch_dims [| m; n |] in\n let out = buffer () (dtype x) out_shape in\n let x_ffi = to_ffi_tensor x' in\n let y_ffi = to_ffi_tensor y' in\n let out_ffi = to_ffi_tensor out in\n caml_matmul x_ffi y_ffi out_ffi;\n out\n\n(* Helper to compute contiguous strides in bytes *)\nlet contiguous_strides shape elem_size =\n let ndim = Array.length shape in\n if ndim = 0 then [||]\n else\n let strides = Array.make ndim 1 in\n for i = ndim - 2 downto 0 do\n strides.(i) <- strides.(i + 1) * shape.(i + 1)\n done;\n Array.map (fun s -> s * elem_size) strides\n\n(* ───── Fourier Transforms Using PocketFFT ───── *)\n\nlet fft (type a b) ?out (x : (a, b) t) ~axes : (a, b) t =\n let x' = materialize x in\n let out_shape = shape x' in\n let out =\n match out with\n | Some o -> o\n | None -> create_tensor x.context x.dtype out_shape\n in\n\n let shape_arr = out_shape in\n let elem_size = Dtype.itemsize x.dtype in\n let strides_in = contiguous_strides out_shape elem_size in\n let strides_out = contiguous_strides out_shape elem_size in\n (* Normalize negative axes *)\n let ndim = Array.length out_shape in\n let axes_arr = Array.map (fun ax -> if ax < 0 then ndim + ax else ax) axes in\n\n (match (x.dtype : (a, b) Dtype.t) with\n | Dtype.Complex64 ->\n Pocketfft.c2c_f32 ~shape:shape_arr ~stride_in:strides_in\n ~stride_out:strides_out ~axes:axes_arr ~forward:true ~fct:1.0\n ~data_in:(Nx_buffer.to_bigarray1 x'.buffer)\n ~data_out:(Nx_buffer.to_bigarray1 out.buffer)\n ~nthreads:1\n | Dtype.Complex128 ->\n Pocketfft.c2c_f64 ~shape:shape_arr ~stride_in:strides_in\n ~stride_out:strides_out ~axes:axes_arr ~forward:true ~fct:1.0\n ~data_in:(Nx_buffer.to_bigarray1 x'.buffer)\n ~data_out:(Nx_buffer.to_bigarray1 out.buffer)\n ~nthreads:1\n | _ -> invalid_arg \"fft: unsupported dtype\");\n\n out\n\nlet ifft (type a b) ?out (x : (a, b) t) ~axes : (a, b) t =\n let x' = materialize x in\n let out_shape = shape x' in\n let out =\n match out with\n | Some o -> o\n | None -> create_tensor x.context x.dtype out_shape\n in\n\n let shape_arr = out_shape in\n let elem_size = Dtype.itemsize x.dtype in\n let strides_in = contiguous_strides out_shape elem_size in\n let strides_out = contiguous_strides out_shape elem_size in\n (* Normalize negative axes *)\n let ndim = Array.length out_shape in\n let axes_arr = Array.map (fun ax -> if ax < 0 then ndim + ax else ax) axes in\n\n (match (x.dtype : (a, b) Dtype.t) with\n | Dtype.Complex64 ->\n Pocketfft.c2c_f32 ~shape:shape_arr ~stride_in:strides_in\n ~stride_out:strides_out ~axes:axes_arr ~forward:false ~fct:1.0\n ~data_in:(Nx_buffer.to_bigarray1 x'.buffer)\n ~data_out:(Nx_buffer.to_bigarray1 out.buffer)\n ~nthreads:1\n | Dtype.Complex128 ->\n Pocketfft.c2c_f64 ~shape:shape_arr ~stride_in:strides_in\n ~stride_out:strides_out ~axes:axes_arr ~forward:false ~fct:1.0\n ~data_in:(Nx_buffer.to_bigarray1 x'.buffer)\n ~data_out:(Nx_buffer.to_bigarray1 out.buffer)\n ~nthreads:1\n | _ -> invalid_arg \"ifft: unsupported dtype\");\n\n out\n\nlet rfft (type a b c d) ?out (x : (a, b) t) ~(dtype : (c, d) Dtype.t) ~axes :\n (c, d) t =\n let x' = materialize x in\n\n (* Calculate output shape for rfft *)\n let in_shape = shape x' in\n let out_shape = Array.copy in_shape in\n let last_axis = Array.length axes - 1 in\n (if last_axis >= 0 then\n let axis_idx =\n if axes.(last_axis) < 0 then Array.length in_shape + axes.(last_axis)\n else axes.(last_axis)\n in\n out_shape.(axis_idx) <- (in_shape.(axis_idx) / 2) + 1);\n\n let out =\n match out with\n | Some o -> o\n | None -> create_tensor x.context dtype out_shape\n in\n\n let strides_in = contiguous_strides in_shape (Dtype.itemsize x.dtype) in\n let strides_out = contiguous_strides out_shape (Dtype.itemsize dtype) in\n\n (* Normalize negative axes *)\n let ndim = Array.length in_shape in\n let axes_normalized =\n Array.map (fun ax -> if ax < 0 then ndim + ax else ax) axes\n in\n\n (match ((x.dtype : (a, b) Dtype.t), (dtype : (c, d) Dtype.t)) with\n | Dtype.Float32, Dtype.Complex64 ->\n Pocketfft.r2c_f32 ~shape_in:in_shape ~stride_in:strides_in\n ~stride_out:strides_out ~axes:axes_normalized ~forward:true ~fct:1.0\n ~data_in:(Nx_buffer.to_bigarray1 x'.buffer)\n ~data_out:(Nx_buffer.to_bigarray1 out.buffer)\n ~nthreads:1\n | Dtype.Float64, Dtype.Complex128 ->\n Pocketfft.r2c_f64 ~shape_in:in_shape ~stride_in:strides_in\n ~stride_out:strides_out ~axes:axes_normalized ~forward:true ~fct:1.0\n ~data_in:(Nx_buffer.to_bigarray1 x'.buffer)\n ~data_out:(Nx_buffer.to_bigarray1 out.buffer)\n ~nthreads:1\n | _ -> invalid_arg \"rfft: unsupported dtype combination\");\n\n out\n\nlet irfft (type a b c d) ?out ?s (x : (a, b) t) ~(dtype : (c, d) Dtype.t) ~axes\n : (c, d) t =\n let x' = materialize x in\n\n (* Calculate output shape for irfft *)\n let in_shape = shape x' in\n let out_shape = Array.copy in_shape in\n let last_axis = Array.length axes - 1 in\n\n (if last_axis >= 0 then\n let axis_idx =\n if axes.(last_axis) < 0 then Array.length in_shape + axes.(last_axis)\n else axes.(last_axis)\n in\n let size =\n match s with\n | None -> (in_shape.(axis_idx) - 1) * 2\n | Some sizes -> sizes.(last_axis)\n in\n out_shape.(axis_idx) <- size);\n\n let out =\n match out with\n | Some o -> o\n | None -> create_tensor x.context dtype out_shape\n in\n\n let strides_in = contiguous_strides in_shape (Dtype.itemsize x.dtype) in\n let strides_out = contiguous_strides out_shape (Dtype.itemsize dtype) in\n\n (* Normalize negative axes *)\n let ndim = Array.length in_shape in\n let axes_normalized =\n Array.map (fun ax -> if ax < 0 then ndim + ax else ax) axes\n in\n\n (match ((x.dtype : (a, b) Dtype.t), (dtype : (c, d) Dtype.t)) with\n | Dtype.Complex64, Dtype.Float32 ->\n Pocketfft.c2r_f32 ~shape_out:out_shape ~stride_in:strides_in\n ~stride_out:strides_out ~axes:axes_normalized ~forward:false ~fct:1.0\n ~data_in:(Nx_buffer.to_bigarray1 x'.buffer)\n ~data_out:(Nx_buffer.to_bigarray1 out.buffer)\n ~nthreads:1\n | Dtype.Complex128, Dtype.Float64 ->\n Pocketfft.c2r_f64 ~shape_out:out_shape ~stride_in:strides_in\n ~stride_out:strides_out ~axes:axes_normalized ~forward:false ~fct:1.0\n ~data_in:(Nx_buffer.to_bigarray1 x'.buffer)\n ~data_out:(Nx_buffer.to_bigarray1 out.buffer)\n ~nthreads:1\n | _ -> invalid_arg \"irfft: unsupported dtype combination\");\n\n out\n\n(* ───── Linear Algebra Operations ───── *)\n\nlet cholesky ~upper x =\n (* Ensure input is materializable *)\n let x' = ensure_materializable x in\n\n (* Create output tensor with same shape and dtype *)\n let out_shape = shape x in\n let out = create_tensor x.context x.dtype out_shape in\n\n (* Convert to FFI tensors *)\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n\n (* Call FFI function *)\n caml_cholesky x_ffi out_ffi upper;\n\n out\n\nlet qr ~reduced x =\n let x' = ensure_materializable x in\n let x_shape = shape x in\n let m = x_shape.(Array.length x_shape - 2) in\n let n = x_shape.(Array.length x_shape - 1) in\n let k = Stdlib.min m n in\n\n (* Calculate Q and R shapes *)\n let q_shape = Array.copy x_shape in\n let r_shape = Array.copy x_shape in\n\n if reduced then (\n (* Reduced QR: Q is m×k, R is k×n *)\n q_shape.(Array.length q_shape - 1) <- k;\n r_shape.(Array.length r_shape - 2) <- k)\n else (\n (* Complete QR: Q is m×m, R is m×n *)\n q_shape.(Array.length q_shape - 1) <- m;\n (* R shape is already m×n from the copy *)\n ());\n\n let q = create_tensor x.context x.dtype q_shape in\n let r = create_tensor x.context x.dtype r_shape in\n\n let x_ffi = to_ffi_tensor x' in\n let q_ffi = to_ffi_tensor q in\n let r_ffi = to_ffi_tensor r in\n\n caml_qr x_ffi q_ffi r_ffi reduced;\n (q, r)\n\nlet svd (type a b) ~full_matrices (x : (a, b) t) :\n (a, b) t * (float, Dtype.float64_elt) t * (a, b) t =\n let x' = ensure_materializable x in\n let x_shape = shape x in\n let m = x_shape.(Array.length x_shape - 2) in\n let n = x_shape.(Array.length x_shape - 1) in\n let k = Stdlib.min m n in\n\n (* Calculate U, S, Vt shapes *)\n let batch_shape = Array.sub x_shape 0 (Array.length x_shape - 2) in\n\n let u_shape =\n Array.append batch_shape (if full_matrices then [| m; m |] else [| m; k |])\n in\n let s_shape = Array.append batch_shape [| k |] in\n let vt_shape =\n Array.append batch_shape (if full_matrices then [| n; n |] else [| k; n |])\n in\n\n let u = create_tensor x.context x.dtype u_shape in\n let s = create_tensor x.context Dtype.Float64 s_shape in\n let vt = create_tensor x.context x.dtype vt_shape in\n\n let x_ffi = to_ffi_tensor x' in\n let u_ffi = to_ffi_tensor u in\n let s_ffi = to_ffi_tensor s in\n let vt_ffi = to_ffi_tensor vt in\n\n caml_svd x_ffi u_ffi s_ffi vt_ffi full_matrices;\n (u, s, vt)\n\nlet eig (type a b) ~vectors (x : (a, b) t) :\n (Complex.t, Dtype.complex64_elt) t\n * (Complex.t, Dtype.complex64_elt) t option =\n let x' = ensure_materializable x in\n let x_shape = shape x in\n let n = x_shape.(Array.length x_shape - 1) in\n\n (* Eigenvalues and eigenvectors are always complex128 *)\n let batch_shape = Array.sub x_shape 0 (Array.length x_shape - 2) in\n let vals_shape = Array.append batch_shape [| n |] in\n let vecs_shape = x_shape in\n\n let vals = create_tensor x.context Dtype.Complex128 vals_shape in\n let vecs =\n if vectors then create_tensor x.context Dtype.Complex128 vecs_shape\n else\n (* Create dummy tensor for C interface *)\n create_tensor x.context Dtype.Complex128 [| 1 |]\n in\n\n let x_ffi = to_ffi_tensor x' in\n let vals_ffi = to_ffi_tensor vals in\n let vecs_ffi = to_ffi_tensor vecs in\n\n caml_eig x_ffi vals_ffi vecs_ffi false vectors;\n if vectors then (vals, Some vecs) else (vals, None)\n\nlet eigh (type a b) ~vectors (x : (a, b) t) :\n (float, Dtype.float64_elt) t * (a, b) t option =\n let x' = ensure_materializable x in\n let x_shape = shape x in\n\n (* For symmetric/hermitian matrices, eigenvalues are always float64 *)\n let batch_shape = Array.sub x_shape 0 (Array.length x_shape - 2) in\n let n = x_shape.(Array.length x_shape - 1) in\n let vals_shape = Array.append batch_shape [| n |] in\n\n let vals = create_tensor x.context Dtype.Float64 vals_shape in\n let vecs =\n if vectors then create_tensor x.context x.dtype x_shape\n else\n (* Create dummy tensor for C interface *)\n create_tensor x.context x.dtype [| 1 |]\n in\n\n let x_ffi = to_ffi_tensor x' in\n let vals_ffi = to_ffi_tensor vals in\n let vecs_ffi = to_ffi_tensor vecs in\n\n caml_eig x_ffi vals_ffi vecs_ffi true vectors;\n if vectors then (vals, Some vecs) else (vals, None)\n\nlet triangular_solve ~upper ~transpose ~unit_diag a b =\n let a' = ensure_materializable a in\n\n (* Handle 1D input b by expanding to 2D *)\n let b_shape = shape b in\n let b_ndim = Array.length b_shape in\n let b_is_1d = b_ndim = 1 in\n\n let b_expanded, out_shape =\n if b_is_1d then\n (* Expand 1D to 2D by adding a trailing dimension *)\n let new_shape = [| b_shape.(0); 1 |] in\n let b_reshaped = reshape b new_shape in\n (b_reshaped, b_shape) (* Keep original shape for output *)\n else (b, shape b)\n in\n\n let b' = ensure_materializable b_expanded in\n\n (* Create output with appropriate shape *)\n let out_shape_expanded = shape b_expanded in\n let out_expanded = create_tensor b.context b.dtype out_shape_expanded in\n\n let a_ffi = to_ffi_tensor a' in\n let b_ffi = to_ffi_tensor b' in\n let out_ffi = to_ffi_tensor out_expanded in\n\n caml_triangular_solve a_ffi b_ffi out_ffi upper transpose unit_diag;\n\n (* Squeeze output back to 1D if input was 1D *)\n if b_is_1d then reshape out_expanded out_shape else out_expanded\n\nlet div x y =\n let dt = dtype x in\n if Dtype.is_int dt || Dtype.is_uint dt then binary_op \"idiv\" caml_idiv x y\n else binary_op \"fdiv\" caml_fdiv x y\n\nlet argmax ~axis ~keepdims x =\n let x' = ensure_materializable x in\n let out_shape = Shape.reduce_output_shape (shape x) [| axis |] keepdims in\n let out = buffer () Dtype.Int32 out_shape in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n caml_argmax x_ffi out_ffi axis keepdims;\n out\n\nlet argmin ~axis ~keepdims x =\n let x' = ensure_materializable x in\n let out_shape = Shape.reduce_output_shape (shape x) [| axis |] keepdims in\n let out = buffer () Dtype.Int32 out_shape in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n caml_argmin x_ffi out_ffi axis keepdims;\n out\n\nlet sort ~axis ~descending x =\n let x' = ensure_materializable x in\n let out = buffer () (dtype x) (shape x) in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n caml_sort x_ffi out_ffi axis descending;\n out\n\nlet argsort ~axis ~descending x =\n let x' = ensure_materializable x in\n let out = buffer () Dtype.Int32 (shape x) in\n let x_ffi = to_ffi_tensor x' in\n let out_ffi = to_ffi_tensor out in\n caml_argsort x_ffi out_ffi axis descending;\n out\n"
},
{
"path": "packages/nx/lib/backend_c/nx_c_binary.c",
"content": "/*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*/\n\n// Binary operations for nx C backend\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"nx_c_shared.h\"\n\n#if defined(_OPENMP)\n#define NX_PARALLEL_THRESHOLD 32768\n#define NX_FOR_EACH_ELEM(total, BODY) \\\n do { \\\n if ((total) >= NX_PARALLEL_THRESHOLD) { \\\n _Pragma(\"omp parallel for simd schedule(static)\") \\\n for (long i = 0; i < (total); ++i) { \\\n BODY; \\\n } \\\n } else { \\\n _Pragma(\"omp simd\") \\\n for (long i = 0; i < (total); ++i) { \\\n BODY; \\\n } \\\n } \\\n } while (0)\n#else\n#define NX_FOR_EACH_ELEM(total, BODY) \\\n do { \\\n for (long i = 0; i < (total); ++i) { \\\n BODY; \\\n } \\\n } while (0)\n#endif\n\n// Type definitions for binary operations\ntypedef void (*binary_op_t)(const ndarray_t *, const ndarray_t *, ndarray_t *);\n\n// Dispatch table for each type\ntypedef struct {\n binary_op_t i8, u8, i16, u16, i32, i64, u32, u64, inat;\n binary_op_t f16, f32, f64;\n binary_op_t c32, c64;\n binary_op_t bf16, bool_, i4, u4, f8e4m3, f8e5m2;\n} binary_op_table;\n\n// Macro to generate all standard type variants for an operation\n// Note: float16, bfloat16, fp8 types need special handling with conversion\n#define GENERATE_BINARY_OP(name, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, int8_t, i8, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, uint8_t, u8, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, int16_t, i16, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, uint16_t, u16, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, int32_t, i32, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, int64_t, i64, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, uint32_t, u32, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, uint64_t, u64, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, intnat, inat, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, float, f32, OP_EXPR) \\\n BINARY_OP_FOR_TYPE(name, double, f64, OP_EXPR)\n\n// Macro to build dispatch table\n#define BUILD_DISPATCH_TABLE(name) \\\n static const binary_op_table name##_table = {.i8 = nx_c_##name##_i8, \\\n .u8 = nx_c_##name##_u8, \\\n .i16 = nx_c_##name##_i16, \\\n .u16 = nx_c_##name##_u16, \\\n .i32 = nx_c_##name##_i32, \\\n .i64 = nx_c_##name##_i64, \\\n .u32 = nx_c_##name##_u32, \\\n .u64 = nx_c_##name##_u64, \\\n .inat = nx_c_##name##_inat, \\\n .f16 = nx_c_##name##_f16, \\\n .f32 = nx_c_##name##_f32, \\\n .f64 = nx_c_##name##_f64, \\\n .c32 = nx_c_##name##_c32, \\\n .c64 = nx_c_##name##_c64, \\\n .bf16 = nx_c_##name##_bf16, \\\n .bool_ = nx_c_##name##_bool_, \\\n .i4 = nx_c_##name##_i4, \\\n .u4 = nx_c_##name##_u4, \\\n .f8e4m3 = nx_c_##name##_f8e4m3, \\\n .f8e5m2 = nx_c_##name##_f8e5m2}\n\n// Helper to iterate over inner dimensions with a kernel function for binary\n// operations\ntypedef void (*kernel_fn)(void *, void *, void *, long, long, long);\n\nstatic inline void iterate_inner_dims(const ndarray_t *x, const ndarray_t *y,\n const ndarray_t *z, long outer_idx,\n kernel_fn kernel, void *x_data,\n void *y_data, void *z_data) {\n if (x->ndim <= 1) {\n kernel(x_data, y_data, z_data, x->offset + outer_idx * x->strides[0],\n y->offset + outer_idx * y->strides[0],\n z->offset + outer_idx * z->strides[0]);\n return;\n }\n\n long x_base = x->offset + outer_idx * x->strides[0];\n long y_base = y->offset + outer_idx * y->strides[0];\n long z_base = z->offset + outer_idx * z->strides[0];\n\n // Create temporary iterator for inner dimensions\n int inner_ndim = x->ndim - 1;\n int coords_stack[MAX_NDIM];\n int *coords = coords_stack;\n bool heap_alloc = false;\n if (inner_ndim > MAX_NDIM) {\n coords = (int *)calloc(inner_ndim, sizeof(int));\n if (!coords) {\n fprintf(stderr, \"nx: iterate_inner_dims: allocation failed\\n\");\n abort();\n }\n heap_alloc = true;\n } else {\n memset(coords_stack, 0, inner_ndim * sizeof(int));\n }\n\n // Iterate over inner dimensions\n bool done = false;\n while (!done) {\n long x_off = x_base;\n long y_off = y_base;\n long z_off = z_base;\n\n for (int i = 0; i < inner_ndim; i++) {\n x_off += coords[i] * x->strides[i + 1];\n y_off += coords[i] * y->strides[i + 1];\n z_off += coords[i] * z->strides[i + 1];\n }\n\n kernel(x_data, y_data, z_data, x_off, y_off, z_off);\n\n // Advance to next position\n done = true;\n for (int i = inner_ndim - 1; i >= 0; i--) {\n coords[i]++;\n if (coords[i] < x->shape[i + 1]) {\n done = false;\n break;\n }\n coords[i] = 0;\n }\n }\n\n if (heap_alloc) free(coords);\n}\n\n// Generic binary operation kernel\n#define BINARY_OP_KERNEL(name, T, suffix, OP) \\\n static inline void nx_c_##name##_##suffix##_kernel(void *x_data, void *y_data, \\\n void *z_data, long x_off, \\\n long y_off, long z_off) { \\\n T *x = (T *)x_data; \\\n T *y = (T *)y_data; \\\n T *z = (T *)z_data; \\\n z[z_off] = OP(x[x_off], y[y_off]); \\\n }\n\n// Generic binary operation implementation\n#define BINARY_OP_IMPL(name, T, suffix) \\\n static void nx_c_##name##_##suffix(const ndarray_t *x, const ndarray_t *y, \\\n ndarray_t *z) { \\\n if (!x || !y || !z) { \\\n fprintf(stderr, \"nx: nx_c_\" #name \"_\" #suffix \": null pointer\\n\"); \\\n abort(); \\\n } \\\n long total = total_elements_safe(x); \\\n if (total == 0) return; \\\n \\\n if (is_fully_contiguous(x, y, z)) { \\\n T *restrict xs = (T *)x->data + x->offset; \\\n T *restrict ys = (T *)y->data + y->offset; \\\n T *restrict zs = (T *)z->data + z->offset; \\\n NX_FOR_EACH_ELEM(total, nx_c_##name##_##suffix##_kernel(xs, ys, zs, i, i, i)); \\\n } else if (x->shape[0] > 1 && total / x->shape[0] > 50) { \\\n _Pragma(\"omp parallel for if(x->shape[0] > 4)\") for (long i = 0; \\\n i < x->shape[0]; \\\n i++) { \\\n iterate_inner_dims(x, y, z, i, nx_c_##name##_##suffix##_kernel, \\\n x->data, y->data, z->data); \\\n } \\\n } else { \\\n nd_iterator_t it; \\\n nd_iterator_init_safe(&it, x, y, z); \\\n do { \\\n long x_off, y_off, z_off; \\\n nd_iterator_get_offsets(&it, &x_off, &y_off, &z_off); \\\n nx_c_##name##_##suffix##_kernel(x->data, y->data, z->data, \\\n x->offset + x_off, y->offset + y_off, \\\n z->offset + z_off); \\\n } while (nd_iterator_next(&it)); \\\n nd_iterator_destroy(&it); \\\n } \\\n }\n\n// Macro to generate both kernel and implementation for an operation\n#define BINARY_OP_FOR_TYPE(name, T, suffix, OP) \\\n BINARY_OP_KERNEL(name, T, suffix, OP) \\\n BINARY_OP_IMPL(name, T, suffix)\n\n// Low-precision float kernel (convert to float for op)\n#define LOW_PREC_OP_KERNEL(name, T, suffix, OP, TO_FLOAT, FROM_FLOAT) \\\n static void nx_c_##name##_##suffix##_kernel(void *x_data, void *y_data, \\\n void *z_data, long x_off, \\\n long y_off, long z_off) { \\\n T *x = (T *)x_data; \\\n T *y = (T *)y_data; \\\n T *z = (T *)z_data; \\\n float a = TO_FLOAT(x[x_off]); \\\n float b = TO_FLOAT(y[y_off]); \\\n z[z_off] = FROM_FLOAT(OP(a, b)); \\\n }\n\n// For floating-point division, no zero check - let IEEE 754 semantics apply\n// (produces inf/-inf/NaN as appropriate)\n\n// For low-precision, use the impl with the special kernel\n#define LOW_PREC_OP_IMPL(name, T, suffix) BINARY_OP_IMPL(name, T, suffix)\n\n// Complex OP for arithmetic - reuse from nx_c_shared.h where possible\n// COMPLEX_ADD and COMPLEX_MUL are defined in nx_c_shared.h\n#define COMPLEX_SUB(x, y) ((x) - (y))\n#define COMPLEX_DIV(x, y) ((x) / (y))\n\n// Helper macros for int4 saturation\n#define CLAMP_I4(x) ((x) < -8 ? -8 : ((x) > 7 ? 7 : (x)))\n#define CLAMP_U4(x) ((x) < 0 ? 0 : ((x) > 15 ? 15 : (x)))\n\n// Special implementation for int4 (packed, unpack/op/pack with saturation)\n#define INT4_OP_IMPL(name, signedness, suffix, OP) \\\n static void nx_c_##name##_##suffix##_kernel(void *x_data, void *y_data, \\\n void *z_data, long x_off, \\\n long y_off, long z_off) { \\\n uint8_t *x = (uint8_t *)x_data; \\\n uint8_t *y = (uint8_t *)y_data; \\\n uint8_t *z = (uint8_t *)z_data; \\\n long byte_off = x_off / 2; \\\n int nib_off = x_off % 2; \\\n int a = nib_off ? (signedness ? (int8_t)(x[byte_off] >> 4) \\\n : (x[byte_off] >> 4) & 0x0F) \\\n : (signedness ? (int8_t)((x[byte_off] & 0x0F) << 4) >> 4 \\\n : x[byte_off] & 0x0F); \\\n int b = nib_off ? (signedness ? (int8_t)(y[byte_off] >> 4) \\\n : (y[byte_off] >> 4) & 0x0F) \\\n : (signedness ? (int8_t)((y[byte_off] & 0x0F) << 4) >> 4 \\\n : y[byte_off] & 0x0F); \\\n int res = OP(a, b); \\\n /* Saturate to 4-bit range */ \\\n res = signedness ? CLAMP_I4(res) : CLAMP_U4(res); \\\n uint8_t nib = (uint8_t)res & 0x0F; \\\n if (nib_off) { \\\n z[byte_off] = (z[byte_off] & 0x0F) | (nib << 4); \\\n } else { \\\n z[byte_off] = (z[byte_off] & 0xF0) | nib; \\\n } \\\n } \\\n static void nx_c_##name##_##suffix(const ndarray_t *x, const ndarray_t *y, \\\n ndarray_t *z) { \\\n if (is_fully_contiguous(x, y, z)) { \\\n long total = total_elements_safe(x); \\\n void *x_data = x->data + x->offset; \\\n void *y_data = y->data + y->offset; \\\n void *z_data = z->data + z->offset; \\\n _Pragma(\"omp parallel for if(total > 10000)\") for (long i = 0; \\\n i < total; i++) { \\\n nx_c_##name##_##suffix##_kernel(x_data, y_data, z_data, i, i, i); \\\n } \\\n } else { \\\n nd_iterator_t it; \\\n nd_iterator_init_safe(&it, x, y, z); \\\n void *x_data = x->data; \\\n void *y_data = y->data; \\\n void *z_data = z->data; \\\n do { \\\n long x_off, y_off, z_off; \\\n nd_iterator_get_offsets(&it, &x_off, &y_off, &z_off); \\\n nx_c_##name##_##suffix##_kernel(x_data, y_data, z_data, \\\n x_off + x->offset, y_off + y->offset, \\\n z_off + z->offset); \\\n } while (nd_iterator_next(&it)); \\\n nd_iterator_destroy(&it); \\\n } \\\n }\n\n// For bool, treat as uint8_t with standard arithmetic\n// Note: Results may exceed 0/1 range (e.g., 1+1=2), stored in uint8_t\n\n// Generate for all ops\n// Addition\n#define ADD_OP(x, y) ((x) + (y))\nGENERATE_BINARY_OP(add, ADD_OP)\n\n// Float16, BFloat16, FP8 variants need conversion\nLOW_PREC_OP_KERNEL(add, uint16_t, f16, ADD_OP, half_to_float, float_to_half)\nLOW_PREC_OP_IMPL(add, uint16_t, f16)\nLOW_PREC_OP_KERNEL(add, caml_ba_bfloat16, bf16, ADD_OP, bfloat16_to_float,\n float_to_bfloat16)\nLOW_PREC_OP_IMPL(add, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(add, caml_ba_fp8_e4m3, f8e4m3, ADD_OP, fp8_e4m3_to_float,\n float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(add, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(add, caml_ba_fp8_e5m2, f8e5m2, ADD_OP, fp8_e5m2_to_float,\n float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(add, caml_ba_fp8_e5m2, f8e5m2)\n\nBINARY_OP_FOR_TYPE(add, complex32, c32, COMPLEX_ADD)\nBINARY_OP_FOR_TYPE(add, complex64, c64, COMPLEX_ADD)\n\nINT4_OP_IMPL(add, 1, i4, ADD_OP)\nINT4_OP_IMPL(add, 0, u4, ADD_OP)\nBINARY_OP_FOR_TYPE(add, caml_ba_bool, bool_, ADD_OP) // Standard arithmetic\nBUILD_DISPATCH_TABLE(add);\n\n// Subtraction\n#define SUB_OP(x, y) ((x) - (y))\nGENERATE_BINARY_OP(sub, SUB_OP)\n\n// Float16, BFloat16, FP8 variants need conversion\nLOW_PREC_OP_KERNEL(sub, uint16_t, f16, SUB_OP, half_to_float, float_to_half)\nLOW_PREC_OP_IMPL(sub, uint16_t, f16)\nLOW_PREC_OP_KERNEL(sub, caml_ba_bfloat16, bf16, SUB_OP, bfloat16_to_float,\n float_to_bfloat16)\nLOW_PREC_OP_IMPL(sub, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(sub, caml_ba_fp8_e4m3, f8e4m3, SUB_OP, fp8_e4m3_to_float,\n float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(sub, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(sub, caml_ba_fp8_e5m2, f8e5m2, SUB_OP, fp8_e5m2_to_float,\n float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(sub, caml_ba_fp8_e5m2, f8e5m2)\n\nBINARY_OP_FOR_TYPE(sub, complex32, c32, COMPLEX_SUB)\nBINARY_OP_FOR_TYPE(sub, complex64, c64, COMPLEX_SUB)\n\nINT4_OP_IMPL(sub, 1, i4, SUB_OP)\nINT4_OP_IMPL(sub, 0, u4, SUB_OP)\nBINARY_OP_FOR_TYPE(sub, caml_ba_bool, bool_,\n SUB_OP) // Standard arithmetic (may wrap)\nBUILD_DISPATCH_TABLE(sub);\n\n// Multiplication\n#define MUL_OP(x, y) ((x) * (y))\nGENERATE_BINARY_OP(mul, MUL_OP)\n\n// Float16, BFloat16, FP8 variants need conversion\nLOW_PREC_OP_KERNEL(mul, uint16_t, f16, MUL_OP, half_to_float, float_to_half)\nLOW_PREC_OP_IMPL(mul, uint16_t, f16)\nLOW_PREC_OP_KERNEL(mul, caml_ba_bfloat16, bf16, MUL_OP, bfloat16_to_float,\n float_to_bfloat16)\nLOW_PREC_OP_IMPL(mul, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(mul, caml_ba_fp8_e4m3, f8e4m3, MUL_OP, fp8_e4m3_to_float,\n float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(mul, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(mul, caml_ba_fp8_e5m2, f8e5m2, MUL_OP, fp8_e5m2_to_float,\n float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(mul, caml_ba_fp8_e5m2, f8e5m2)\n\nBINARY_OP_FOR_TYPE(mul, complex32, c32, COMPLEX_MUL)\nBINARY_OP_FOR_TYPE(mul, complex64, c64, COMPLEX_MUL)\n\nINT4_OP_IMPL(mul, 1, i4, MUL_OP)\nINT4_OP_IMPL(mul, 0, u4, MUL_OP)\nBINARY_OP_FOR_TYPE(mul, caml_ba_bool, bool_, MUL_OP) // Standard arithmetic\nBUILD_DISPATCH_TABLE(mul);\n\n// Integer division - truncates and checks for zero\n#define INT_DIV_OP(x, y) \\\n ((y) == 0 ? (fprintf(stderr, \"nx: division by zero\\n\"), abort(), (x)) : ((x) / (y)))\n\n// Integer types\nBINARY_OP_FOR_TYPE(idiv, int8_t, i8, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, uint8_t, u8, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, int16_t, i16, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, uint16_t, u16, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, int32_t, i32, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, int64_t, i64, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, uint32_t, u32, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, uint64_t, u64, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, intnat, inat, INT_DIV_OP)\n\n// For float types, idiv truncates the result\n#define FLOAT_IDIV_OP(x, y) (trunc((x) / (y)))\nBINARY_OP_FOR_TYPE(idiv, float, f32, FLOAT_IDIV_OP)\nBINARY_OP_FOR_TYPE(idiv, double, f64, FLOAT_IDIV_OP)\n\nLOW_PREC_OP_KERNEL(idiv, uint16_t, f16, FLOAT_IDIV_OP, half_to_float,\n float_to_half)\nLOW_PREC_OP_IMPL(idiv, uint16_t, f16)\nLOW_PREC_OP_KERNEL(idiv, caml_ba_bfloat16, bf16, FLOAT_IDIV_OP,\n bfloat16_to_float, float_to_bfloat16)\nLOW_PREC_OP_IMPL(idiv, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(idiv, caml_ba_fp8_e4m3, f8e4m3, FLOAT_IDIV_OP,\n fp8_e4m3_to_float, float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(idiv, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(idiv, caml_ba_fp8_e5m2, f8e5m2, FLOAT_IDIV_OP,\n fp8_e5m2_to_float, float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(idiv, caml_ba_fp8_e5m2, f8e5m2)\n\n// Complex idiv also truncates both real and imaginary parts\nstatic void nx_c_idiv_c32_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n complex32 *x = (complex32 *)x_data;\n complex32 *y = (complex32 *)y_data;\n complex32 *z = (complex32 *)z_data;\n complex32 res = x[x_off] / y[y_off];\n z[z_off] = truncf(crealf(res)) + I * truncf(cimagf(res));\n}\n\nstatic void nx_c_idiv_c64_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n complex64 *x = (complex64 *)x_data;\n complex64 *y = (complex64 *)y_data;\n complex64 *z = (complex64 *)z_data;\n complex64 res = x[x_off] / y[y_off];\n z[z_off] = trunc(creal(res)) + I * trunc(cimag(res));\n}\n\nBINARY_OP_IMPL(idiv, complex32, c32)\nBINARY_OP_IMPL(idiv, complex64, c64)\n\n\nINT4_OP_IMPL(idiv, 1, i4, INT_DIV_OP)\nINT4_OP_IMPL(idiv, 0, u4, INT_DIV_OP)\nBINARY_OP_FOR_TYPE(idiv, caml_ba_bool, bool_, INT_DIV_OP)\nBUILD_DISPATCH_TABLE(idiv);\n\n// Floating-point division - follows IEEE 754 (inf/NaN for division by zero)\n#define FLOAT_DIV_OP(x, y) ((x) / (y))\n\n// Integer types converted to float for fdiv\nBINARY_OP_FOR_TYPE(fdiv, int8_t, i8, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, uint8_t, u8, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, int16_t, i16, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, uint16_t, u16, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, int32_t, i32, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, int64_t, i64, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, uint32_t, u32, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, uint64_t, u64, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, intnat, inat, FLOAT_DIV_OP)\n\n// Floating-point types use IEEE 754 semantics\nBINARY_OP_FOR_TYPE(fdiv, float, f32, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, double, f64, FLOAT_DIV_OP)\n\nLOW_PREC_OP_KERNEL(fdiv, uint16_t, f16, FLOAT_DIV_OP, half_to_float,\n float_to_half)\nLOW_PREC_OP_IMPL(fdiv, uint16_t, f16)\nLOW_PREC_OP_KERNEL(fdiv, caml_ba_bfloat16, bf16, FLOAT_DIV_OP,\n bfloat16_to_float, float_to_bfloat16)\nLOW_PREC_OP_IMPL(fdiv, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(fdiv, caml_ba_fp8_e4m3, f8e4m3, FLOAT_DIV_OP,\n fp8_e4m3_to_float, float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(fdiv, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(fdiv, caml_ba_fp8_e5m2, f8e5m2, FLOAT_DIV_OP,\n fp8_e5m2_to_float, float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(fdiv, caml_ba_fp8_e5m2, f8e5m2)\n\nBINARY_OP_FOR_TYPE(fdiv, complex32, c32, COMPLEX_DIV)\nBINARY_OP_FOR_TYPE(fdiv, complex64, c64, COMPLEX_DIV)\n\nINT4_OP_IMPL(fdiv, 1, i4, FLOAT_DIV_OP)\nINT4_OP_IMPL(fdiv, 0, u4, FLOAT_DIV_OP)\nBINARY_OP_FOR_TYPE(fdiv, caml_ba_bool, bool_, FLOAT_DIV_OP)\nBUILD_DISPATCH_TABLE(fdiv);\n\n// Max/Min with special complex\n#define MAX_OP(x, y) ((x) > (y) ? (x) : (y))\nGENERATE_BINARY_OP(max, MAX_OP)\n\n// Float16, BFloat16, FP8 variants need conversion\nLOW_PREC_OP_KERNEL(max, uint16_t, f16, MAX_OP, half_to_float, float_to_half)\nLOW_PREC_OP_IMPL(max, uint16_t, f16)\nLOW_PREC_OP_KERNEL(max, caml_ba_bfloat16, bf16, MAX_OP, bfloat16_to_float,\n float_to_bfloat16)\nLOW_PREC_OP_IMPL(max, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(max, caml_ba_fp8_e4m3, f8e4m3, MAX_OP, fp8_e4m3_to_float,\n float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(max, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(max, caml_ba_fp8_e5m2, f8e5m2, MAX_OP, fp8_e5m2_to_float,\n float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(max, caml_ba_fp8_e5m2, f8e5m2)\n\nBINARY_OP_FOR_TYPE(max, complex32, c32, complex_max)\nBINARY_OP_FOR_TYPE(max, complex64, c64, complex64_max)\nINT4_OP_IMPL(max, 1, i4, MAX_OP)\nINT4_OP_IMPL(max, 0, u4, MAX_OP)\nBINARY_OP_FOR_TYPE(max, caml_ba_bool, bool_, MAX_OP) // Standard comparison\nBUILD_DISPATCH_TABLE(max);\n\n#define MIN_OP(x, y) ((x) < (y) ? (x) : (y))\nGENERATE_BINARY_OP(min, MIN_OP)\n\n// Float16, BFloat16, FP8 variants need conversion\nLOW_PREC_OP_KERNEL(min, uint16_t, f16, MIN_OP, half_to_float, float_to_half)\nLOW_PREC_OP_IMPL(min, uint16_t, f16)\nLOW_PREC_OP_KERNEL(min, caml_ba_bfloat16, bf16, MIN_OP, bfloat16_to_float,\n float_to_bfloat16)\nLOW_PREC_OP_IMPL(min, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(min, caml_ba_fp8_e4m3, f8e4m3, MIN_OP, fp8_e4m3_to_float,\n float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(min, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(min, caml_ba_fp8_e5m2, f8e5m2, MIN_OP, fp8_e5m2_to_float,\n float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(min, caml_ba_fp8_e5m2, f8e5m2)\n\nBINARY_OP_FOR_TYPE(min, complex32, c32, complex_min)\nBINARY_OP_FOR_TYPE(min, complex64, c64, complex64_min)\nINT4_OP_IMPL(min, 1, i4, MIN_OP)\nINT4_OP_IMPL(min, 0, u4, MIN_OP)\nBINARY_OP_FOR_TYPE(min, caml_ba_bool, bool_, MIN_OP) // Standard comparison\nBUILD_DISPATCH_TABLE(min);\n\n// =========== MODULO ===========\n#define MOD_OP(x, y) \\\n ((y) == 0 ? (fprintf(stderr, \"nx: modulo by zero\\n\"), abort(), 0) : ((x) % (y)))\n#define FMOD_OP(x, y) (fmod((x), (y)))\n\n// Integer modulo\nBINARY_OP_FOR_TYPE(mod, int8_t, i8, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, uint8_t, u8, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, int16_t, i16, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, uint16_t, u16, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, int32_t, i32, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, int64_t, i64, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, uint32_t, u32, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, uint64_t, u64, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, intnat, inat, MOD_OP)\n\n// Float modulo uses fmod\nBINARY_OP_FOR_TYPE(mod, float, f32, FMOD_OP)\nBINARY_OP_FOR_TYPE(mod, double, f64, FMOD_OP)\nLOW_PREC_OP_KERNEL(mod, uint16_t, f16, FMOD_OP, half_to_float, float_to_half)\nLOW_PREC_OP_IMPL(mod, uint16_t, f16)\nLOW_PREC_OP_KERNEL(mod, caml_ba_bfloat16, bf16, FMOD_OP, bfloat16_to_float,\n float_to_bfloat16)\nLOW_PREC_OP_IMPL(mod, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(mod, caml_ba_fp8_e4m3, f8e4m3, FMOD_OP, fp8_e4m3_to_float,\n float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(mod, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(mod, caml_ba_fp8_e5m2, f8e5m2, FMOD_OP, fp8_e5m2_to_float,\n float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(mod, caml_ba_fp8_e5m2, f8e5m2)\n\n// Complex modulo not well-defined\n\nINT4_OP_IMPL(mod, 1, i4, MOD_OP)\nINT4_OP_IMPL(mod, 0, u4, MOD_OP)\nBINARY_OP_FOR_TYPE(mod, caml_ba_bool, bool_, MOD_OP)\n\n// Build dispatch table with NULL for unsupported complex types\nstatic const binary_op_table mod_table = {.i8 = nx_c_mod_i8,\n .u8 = nx_c_mod_u8,\n .i16 = nx_c_mod_i16,\n .u16 = nx_c_mod_u16,\n .i32 = nx_c_mod_i32,\n .i64 = nx_c_mod_i64,\n .u32 = nx_c_mod_u32,\n .u64 = nx_c_mod_u64,\n .inat = nx_c_mod_inat,\n .f16 = nx_c_mod_f16,\n .f32 = nx_c_mod_f32,\n .f64 = nx_c_mod_f64,\n .c32 = NULL,\n .c64 = NULL,\n .bf16 = nx_c_mod_bf16,\n .bool_ = nx_c_mod_bool_,\n .i4 = nx_c_mod_i4,\n .u4 = nx_c_mod_u4,\n .f8e4m3 = nx_c_mod_f8e4m3,\n .f8e5m2 = nx_c_mod_f8e5m2};\n\n// =========== POWER ===========\n#define POW_OP(x, y) (pow((double)(x), (double)(y)))\n#define FPOW_OP(x, y) (powf((x), (y)))\n#define DPOW_OP(x, y) (pow((x), (y)))\n\n// All types use pow, converting to appropriate precision\nBINARY_OP_FOR_TYPE(pow, int8_t, i8, POW_OP)\nBINARY_OP_FOR_TYPE(pow, uint8_t, u8, POW_OP)\nBINARY_OP_FOR_TYPE(pow, int16_t, i16, POW_OP)\nBINARY_OP_FOR_TYPE(pow, uint16_t, u16, POW_OP)\nBINARY_OP_FOR_TYPE(pow, int32_t, i32, POW_OP)\nBINARY_OP_FOR_TYPE(pow, int64_t, i64, POW_OP)\nBINARY_OP_FOR_TYPE(pow, uint32_t, u32, POW_OP)\nBINARY_OP_FOR_TYPE(pow, uint64_t, u64, POW_OP)\nBINARY_OP_FOR_TYPE(pow, intnat, inat, POW_OP)\n\nBINARY_OP_FOR_TYPE(pow, float, f32, FPOW_OP)\nBINARY_OP_FOR_TYPE(pow, double, f64, DPOW_OP)\nLOW_PREC_OP_KERNEL(pow, uint16_t, f16, FPOW_OP, half_to_float, float_to_half)\nLOW_PREC_OP_IMPL(pow, uint16_t, f16)\nLOW_PREC_OP_KERNEL(pow, caml_ba_bfloat16, bf16, FPOW_OP, bfloat16_to_float,\n float_to_bfloat16)\nLOW_PREC_OP_IMPL(pow, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(pow, caml_ba_fp8_e4m3, f8e4m3, FPOW_OP, fp8_e4m3_to_float,\n float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(pow, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(pow, caml_ba_fp8_e5m2, f8e5m2, FPOW_OP, fp8_e5m2_to_float,\n float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(pow, caml_ba_fp8_e5m2, f8e5m2)\n\n// Complex power using cpow\n#define CPOW32_OP(x, y) (cpowf((x), (y)))\n#define CPOW64_OP(x, y) (cpow((x), (y)))\nBINARY_OP_FOR_TYPE(pow, complex32, c32, CPOW32_OP)\nBINARY_OP_FOR_TYPE(pow, complex64, c64, CPOW64_OP)\n\n\nINT4_OP_IMPL(pow, 1, i4, POW_OP)\nINT4_OP_IMPL(pow, 0, u4, POW_OP)\nBINARY_OP_FOR_TYPE(pow, caml_ba_bool, bool_, POW_OP)\nBUILD_DISPATCH_TABLE(pow);\n\n// =========== ATAN2 ===========\n#define ATAN2F_OP(x, y) (atan2f((x), (y)))\n#define ATAN2D_OP(x, y) (atan2((x), (y)))\n\nBINARY_OP_FOR_TYPE(atan2, float, f32, ATAN2F_OP)\nBINARY_OP_FOR_TYPE(atan2, double, f64, ATAN2D_OP)\n\nLOW_PREC_OP_KERNEL(atan2, uint16_t, f16, ATAN2F_OP, half_to_float, float_to_half)\nLOW_PREC_OP_IMPL(atan2, uint16_t, f16)\nLOW_PREC_OP_KERNEL(atan2, caml_ba_bfloat16, bf16, ATAN2F_OP, bfloat16_to_float,\n float_to_bfloat16)\nLOW_PREC_OP_IMPL(atan2, caml_ba_bfloat16, bf16)\nLOW_PREC_OP_KERNEL(atan2, caml_ba_fp8_e4m3, f8e4m3, ATAN2F_OP,\n fp8_e4m3_to_float, float_to_fp8_e4m3)\nLOW_PREC_OP_IMPL(atan2, caml_ba_fp8_e4m3, f8e4m3)\nLOW_PREC_OP_KERNEL(atan2, caml_ba_fp8_e5m2, f8e5m2, ATAN2F_OP,\n fp8_e5m2_to_float, float_to_fp8_e5m2)\nLOW_PREC_OP_IMPL(atan2, caml_ba_fp8_e5m2, f8e5m2)\n\nstatic const binary_op_table atan2_table = {.i8 = NULL,\n .u8 = NULL,\n .i16 = NULL,\n .u16 = NULL,\n .i32 = NULL,\n .i64 = NULL,\n .u32 = NULL,\n .u64 = NULL,\n .inat = NULL,\n .f16 = nx_c_atan2_f16,\n .f32 = nx_c_atan2_f32,\n .f64 = nx_c_atan2_f64,\n .c32 = NULL,\n .c64 = NULL,\n .bf16 = nx_c_atan2_bf16,\n .bool_ = NULL,\n .i4 = NULL,\n .u4 = NULL,\n .f8e4m3 = nx_c_atan2_f8e4m3,\n .f8e5m2 = nx_c_atan2_f8e5m2};\n\n// =========== COMPARISON - LESS THAN ===========\n#define CMPLT_OP(x, y) ((x) < (y) ? 1 : 0)\n\n// Comparison operations that output uint8\n#define COMPARISON_OP_FOR_TYPE(name, T, suffix, OP) \\\n static void nx_c_##name##_##suffix##_kernel(void *x_data, void *y_data, \\\n void *z_data, long x_off, \\\n long y_off, long z_off) { \\\n T *x = (T *)x_data; \\\n T *y = (T *)y_data; \\\n uint8_t *z = (uint8_t *)z_data; \\\n z[z_off] = OP(x[x_off], y[y_off]); \\\n } \\\n static void nx_c_##name##_##suffix(const ndarray_t *x, const ndarray_t *y, \\\n ndarray_t *z) { \\\n if (!x || !y || !z) { \\\n fprintf(stderr, \"nx: nx_c_\" #name \"_\" #suffix \": null pointer\\n\"); \\\n abort(); \\\n } \\\n long total = total_elements_safe(x); \\\n if (total == 0) return; \\\n \\\n if (is_fully_contiguous(x, y, z)) { \\\n _Pragma(\"omp parallel for simd if(total > 1000)\") for (long i = 0; \\\n i < total; i++) { \\\n nx_c_##name##_##suffix##_kernel(x->data, y->data, z->data, \\\n x->offset + i, y->offset + i, \\\n z->offset + i); \\\n } \\\n } else if (x->shape[0] > 1 && total / x->shape[0] > 50) { \\\n _Pragma(\"omp parallel for if(x->shape[0] > 4)\") for (long i = 0; \\\n i < x->shape[0]; \\\n i++) { \\\n iterate_inner_dims(x, y, z, i, nx_c_##name##_##suffix##_kernel, \\\n x->data, y->data, z->data); \\\n } \\\n } else { \\\n nd_iterator_t it; \\\n nd_iterator_init_safe(&it, x, y, z); \\\n do { \\\n long x_off, y_off, z_off; \\\n nd_iterator_get_offsets(&it, &x_off, &y_off, &z_off); \\\n nx_c_##name##_##suffix##_kernel(x->data, y->data, z->data, \\\n x->offset + x_off, y->offset + y_off, \\\n z->offset + z_off); \\\n } while (nd_iterator_next(&it)); \\\n nd_iterator_destroy(&it); \\\n } \\\n }\n\nCOMPARISON_OP_FOR_TYPE(cmplt, int8_t, i8, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, uint8_t, u8, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, int16_t, i16, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, uint16_t, u16, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, int32_t, i32, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, int64_t, i64, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, uint32_t, u32, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, uint64_t, u64, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, intnat, inat, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, float, f32, CMPLT_OP)\nCOMPARISON_OP_FOR_TYPE(cmplt, double, f64, CMPLT_OP)\n\n// Low precision comparisons\nstatic void nx_c_cmplt_f16_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n uint16_t *x = (uint16_t *)x_data;\n uint16_t *y = (uint16_t *)y_data;\n bool *z = (bool *)z_data;\n float a = half_to_float(x[x_off]);\n float b = half_to_float(y[y_off]);\n z[z_off] = a < b ? 1 : 0;\n}\nBINARY_OP_IMPL(cmplt, uint16_t, f16)\n\n// Similar for other low-precision types\n#define LOW_PREC_CMP_KERNEL(name, T, suffix, OP, TO_FLOAT) \\\n static void nx_c_##name##_##suffix##_kernel(void *x_data, void *y_data, \\\n void *z_data, long x_off, \\\n long y_off, long z_off) { \\\n T *x = (T *)x_data; \\\n T *y = (T *)y_data; \\\n bool *z = (bool *)z_data; \\\n float a = TO_FLOAT(x[x_off]); \\\n float b = TO_FLOAT(y[y_off]); \\\n z[z_off] = OP(a, b); \\\n } \\\n BINARY_OP_IMPL(name, T, suffix)\n\nLOW_PREC_CMP_KERNEL(cmplt, caml_ba_bfloat16, bf16, CMPLT_OP, bfloat16_to_float)\nLOW_PREC_CMP_KERNEL(cmplt, caml_ba_fp8_e4m3, f8e4m3, CMPLT_OP,\n fp8_e4m3_to_float)\nLOW_PREC_CMP_KERNEL(cmplt, caml_ba_fp8_e5m2, f8e5m2, CMPLT_OP,\n fp8_e5m2_to_float)\n\n// Complex comparison not well-defined\n\n// Int4 comparison implementation - unpacks 4-bit values and outputs uint8\n#define INT4_COMPARISON_OP_IMPL(name, signedness, suffix, OP) \\\n static void nx_c_##name##_##suffix##_kernel(void *x_data, void *y_data, \\\n void *z_data, long x_off, \\\n long y_off, long z_off) { \\\n uint8_t *x = (uint8_t *)x_data; \\\n uint8_t *y = (uint8_t *)y_data; \\\n bool *z = (bool *)z_data; \\\n /* Unpack x value */ \\\n long x_byte_off = x_off / 2; \\\n int x_nib_off = x_off % 2; \\\n int a = x_nib_off ? (signedness ? (int8_t)(x[x_byte_off] >> 4) \\\n : (x[x_byte_off] >> 4) & 0x0F) \\\n : (signedness ? (int8_t)((x[x_byte_off] & 0x0F) << 4) >> 4 \\\n : x[x_byte_off] & 0x0F); \\\n /* Unpack y value */ \\\n long y_byte_off = y_off / 2; \\\n int y_nib_off = y_off % 2; \\\n int b = y_nib_off ? (signedness ? (int8_t)(y[y_byte_off] >> 4) \\\n : (y[y_byte_off] >> 4) & 0x0F) \\\n : (signedness ? (int8_t)((y[y_byte_off] & 0x0F) << 4) >> 4 \\\n : y[y_byte_off] & 0x0F); \\\n /* Store comparison result as uint8 (0 or 1) */ \\\n z[z_off] = OP(a, b); \\\n } \\\n static void nx_c_##name##_##suffix(const ndarray_t *x, const ndarray_t *y, \\\n ndarray_t *z) { \\\n if (is_fully_contiguous(x, y, z)) { \\\n long total = total_elements_safe(x); \\\n void *x_data = x->data + x->offset; \\\n void *y_data = y->data + y->offset; \\\n void *z_data = z->data + z->offset; \\\n _Pragma(\"omp parallel for if(total > 10000)\") for (long i = 0; \\\n i < total; i++) { \\\n nx_c_##name##_##suffix##_kernel(x_data, y_data, z_data, i, i, i); \\\n } \\\n } else { \\\n nd_iterator_t it; \\\n nd_iterator_init_safe(&it, x, y, z); \\\n void *x_data = x->data; \\\n void *y_data = y->data; \\\n void *z_data = z->data; \\\n do { \\\n long x_off, y_off, z_off; \\\n nd_iterator_get_offsets(&it, &x_off, &y_off, &z_off); \\\n nx_c_##name##_##suffix##_kernel(x_data, y_data, z_data, \\\n x_off + x->offset, y_off + y->offset, \\\n z_off + z->offset); \\\n } while (nd_iterator_next(&it)); \\\n nd_iterator_destroy(&it); \\\n } \\\n }\n\n// Define comparison operators\n#define CMPGT_OP(x, y) ((x) > (y) ? true : false)\n#define CMPLE_OP(x, y) ((x) <= (y) ? true : false)\n#define CMPGE_OP(x, y) ((x) >= (y) ? true : false)\n#define CMPEQ_OP(x, y) ((x) == (y) ? true : false)\n#define CMPNE_OP(x, y) ((x) != (y) ? true : false)\n\n// Generate int4/uint4 comparison operations\nINT4_COMPARISON_OP_IMPL(cmplt, 1, i4, CMPLT_OP)\nINT4_COMPARISON_OP_IMPL(cmplt, 0, u4, CMPLT_OP)\nINT4_COMPARISON_OP_IMPL(cmpgt, 1, i4, CMPGT_OP)\nINT4_COMPARISON_OP_IMPL(cmpgt, 0, u4, CMPGT_OP)\nINT4_COMPARISON_OP_IMPL(cmple, 1, i4, CMPLE_OP)\nINT4_COMPARISON_OP_IMPL(cmple, 0, u4, CMPLE_OP)\nINT4_COMPARISON_OP_IMPL(cmpge, 1, i4, CMPGE_OP)\nINT4_COMPARISON_OP_IMPL(cmpge, 0, u4, CMPGE_OP)\nINT4_COMPARISON_OP_IMPL(cmpeq, 1, i4, CMPEQ_OP)\nINT4_COMPARISON_OP_IMPL(cmpeq, 0, u4, CMPEQ_OP)\nINT4_COMPARISON_OP_IMPL(cmpne, 1, i4, CMPNE_OP)\nINT4_COMPARISON_OP_IMPL(cmpne, 0, u4, CMPNE_OP)\n\nCOMPARISON_OP_FOR_TYPE(cmplt, caml_ba_bool, bool_, CMPLT_OP)\n\n// Build dispatch table with NULL for unsupported complex and int4/uint4 types\nstatic const binary_op_table cmplt_table = {.i8 = nx_c_cmplt_i8,\n .u8 = nx_c_cmplt_u8,\n .i16 = nx_c_cmplt_i16,\n .u16 = nx_c_cmplt_u16,\n .i32 = nx_c_cmplt_i32,\n .i64 = nx_c_cmplt_i64,\n .u32 = nx_c_cmplt_u32,\n .u64 = nx_c_cmplt_u64,\n .inat = nx_c_cmplt_inat,\n .f16 = nx_c_cmplt_f16,\n .f32 = nx_c_cmplt_f32,\n .f64 = nx_c_cmplt_f64,\n .c32 = NULL,\n .c64 = NULL,\n .bf16 = nx_c_cmplt_bf16,\n .bool_ = nx_c_cmplt_bool_,\n .i4 = nx_c_cmplt_i4,\n .u4 = nx_c_cmplt_u4,\n .f8e4m3 = nx_c_cmplt_f8e4m3,\n .f8e5m2 = nx_c_cmplt_f8e5m2};\n\n// =========== COMPARISON - NOT EQUAL ===========\nCOMPARISON_OP_FOR_TYPE(cmpne, int8_t, i8, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, uint8_t, u8, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, int16_t, i16, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, uint16_t, u16, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, int32_t, i32, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, int64_t, i64, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, uint32_t, u32, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, uint64_t, u64, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, intnat, inat, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, float, f32, CMPNE_OP)\nCOMPARISON_OP_FOR_TYPE(cmpne, double, f64, CMPNE_OP)\n\n// Low precision\nstatic void nx_c_cmpne_f16_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n uint16_t *x = (uint16_t *)x_data;\n uint16_t *y = (uint16_t *)y_data;\n uint8_t *z = (uint8_t *)z_data;\n float a = half_to_float(x[x_off]);\n float b = half_to_float(y[y_off]);\n z[z_off] = a != b ? 1 : 0;\n}\nBINARY_OP_IMPL(cmpne, uint16_t, f16)\n\nLOW_PREC_CMP_KERNEL(cmpne, caml_ba_bfloat16, bf16, CMPNE_OP, bfloat16_to_float)\nLOW_PREC_CMP_KERNEL(cmpne, caml_ba_fp8_e4m3, f8e4m3, CMPNE_OP,\n fp8_e4m3_to_float)\nLOW_PREC_CMP_KERNEL(cmpne, caml_ba_fp8_e5m2, f8e5m2, CMPNE_OP,\n fp8_e5m2_to_float)\n\n// Complex comparison for equality\nstatic void nx_c_cmpne_c32_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n complex32 *x = (complex32 *)x_data;\n complex32 *y = (complex32 *)y_data;\n uint8_t *z = (uint8_t *)z_data;\n z[z_off] = (x[x_off] != y[y_off]) ? 1 : 0;\n}\nBINARY_OP_IMPL(cmpne, complex32, c32)\n\nstatic void nx_c_cmpne_c64_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n complex64 *x = (complex64 *)x_data;\n complex64 *y = (complex64 *)y_data;\n uint8_t *z = (uint8_t *)z_data;\n z[z_off] = (x[x_off] != y[y_off]) ? 1 : 0;\n}\nBINARY_OP_IMPL(cmpne, complex64, c64)\n\n// Int4 comparison not yet implemented\n\nCOMPARISON_OP_FOR_TYPE(cmpne, caml_ba_bool, bool_, CMPNE_OP)\n\n// Build dispatch table with NULL for unsupported int4/uint4 types\nstatic const binary_op_table cmpne_table = {.i8 = nx_c_cmpne_i8,\n .u8 = nx_c_cmpne_u8,\n .i16 = nx_c_cmpne_i16,\n .u16 = nx_c_cmpne_u16,\n .i32 = nx_c_cmpne_i32,\n .i64 = nx_c_cmpne_i64,\n .u32 = nx_c_cmpne_u32,\n .u64 = nx_c_cmpne_u64,\n .inat = nx_c_cmpne_inat,\n .f16 = nx_c_cmpne_f16,\n .f32 = nx_c_cmpne_f32,\n .f64 = nx_c_cmpne_f64,\n .c32 = nx_c_cmpne_c32,\n .c64 = nx_c_cmpne_c64,\n .bf16 = nx_c_cmpne_bf16,\n .bool_ = nx_c_cmpne_bool_,\n .i4 = nx_c_cmpne_i4,\n .u4 = nx_c_cmpne_u4,\n .f8e4m3 = nx_c_cmpne_f8e4m3,\n .f8e5m2 = nx_c_cmpne_f8e5m2};\n\n// =========== COMPARISON - EQUAL ===========\nCOMPARISON_OP_FOR_TYPE(cmpeq, int8_t, i8, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, uint8_t, u8, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, int16_t, i16, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, uint16_t, u16, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, int32_t, i32, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, int64_t, i64, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, uint32_t, u32, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, uint64_t, u64, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, intnat, inat, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, float, f32, CMPEQ_OP)\nCOMPARISON_OP_FOR_TYPE(cmpeq, double, f64, CMPEQ_OP)\n\n// Low precision\nstatic void nx_c_cmpeq_f16_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n uint16_t *x = (uint16_t *)x_data;\n uint16_t *y = (uint16_t *)y_data;\n uint8_t *z = (uint8_t *)z_data;\n float a = half_to_float(x[x_off]);\n float b = half_to_float(y[y_off]);\n z[z_off] = a == b ? 1 : 0;\n}\nBINARY_OP_IMPL(cmpeq, uint16_t, f16)\n\nLOW_PREC_CMP_KERNEL(cmpeq, caml_ba_bfloat16, bf16, CMPEQ_OP, bfloat16_to_float)\nLOW_PREC_CMP_KERNEL(cmpeq, caml_ba_fp8_e4m3, f8e4m3, CMPEQ_OP,\n fp8_e4m3_to_float)\nLOW_PREC_CMP_KERNEL(cmpeq, caml_ba_fp8_e5m2, f8e5m2, CMPEQ_OP,\n fp8_e5m2_to_float)\n\n// Complex comparison for equality\nstatic void nx_c_cmpeq_c32_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n complex32 *x = (complex32 *)x_data;\n complex32 *y = (complex32 *)y_data;\n uint8_t *z = (uint8_t *)z_data;\n z[z_off] = (x[x_off] == y[y_off]) ? 1 : 0;\n}\nBINARY_OP_IMPL(cmpeq, complex32, c32)\n\nstatic void nx_c_cmpeq_c64_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n complex64 *x = (complex64 *)x_data;\n complex64 *y = (complex64 *)y_data;\n uint8_t *z = (uint8_t *)z_data;\n z[z_off] = (x[x_off] == y[y_off]) ? 1 : 0;\n}\nBINARY_OP_IMPL(cmpeq, complex64, c64)\n\nCOMPARISON_OP_FOR_TYPE(cmpeq, caml_ba_bool, bool_, CMPEQ_OP)\n\nstatic const binary_op_table cmpeq_table = {.i8 = nx_c_cmpeq_i8,\n .u8 = nx_c_cmpeq_u8,\n .i16 = nx_c_cmpeq_i16,\n .u16 = nx_c_cmpeq_u16,\n .i32 = nx_c_cmpeq_i32,\n .i64 = nx_c_cmpeq_i64,\n .u32 = nx_c_cmpeq_u32,\n .u64 = nx_c_cmpeq_u64,\n .inat = nx_c_cmpeq_inat,\n .f16 = nx_c_cmpeq_f16,\n .f32 = nx_c_cmpeq_f32,\n .f64 = nx_c_cmpeq_f64,\n .c32 = nx_c_cmpeq_c32,\n .c64 = nx_c_cmpeq_c64,\n .bf16 = nx_c_cmpeq_bf16,\n .bool_ = nx_c_cmpeq_bool_,\n .i4 = nx_c_cmpeq_i4,\n .u4 = nx_c_cmpeq_u4,\n .f8e4m3 = nx_c_cmpeq_f8e4m3,\n .f8e5m2 = nx_c_cmpeq_f8e5m2};\n\n// =========== COMPARISON - LESS THAN OR EQUAL ===========\nCOMPARISON_OP_FOR_TYPE(cmple, int8_t, i8, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, uint8_t, u8, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, int16_t, i16, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, uint16_t, u16, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, int32_t, i32, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, int64_t, i64, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, uint32_t, u32, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, uint64_t, u64, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, intnat, inat, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, float, f32, CMPLE_OP)\nCOMPARISON_OP_FOR_TYPE(cmple, double, f64, CMPLE_OP)\n\n// Low precision comparisons\nstatic void nx_c_cmple_f16_kernel(void *x_data, void *y_data, void *z_data,\n long x_off, long y_off, long z_off) {\n uint16_t *x = (uint16_t *)x_data;\n uint16_t *y = (uint16_t *)y_data;\n uint8_t *z = (uint8_t *)z_data;\n float a = half_to_float(x[x_off]);\n float b = half_to_float(y[y_off]);\n z[z_off] = a <= b ? 1 : 0;\n}\nBINARY_OP_IMPL(cmple, uint16_t, f16)\n\nLOW_PREC_CMP_KERNEL(cmple, caml_ba_bfloat16, bf16, CMPLE_OP, bfloat16_to_float)\nLOW_PREC_CMP_KERNEL(cmple, caml_ba_fp8_e4m3, f8e4m3, CMPLE_OP,\n fp8_e4m3_to_float)\nLOW_PREC_CMP_KERNEL(cmple, caml_ba_fp8_e5m2, f8e5m2, CMPLE_OP,\n fp8_e5m2_to_float)\n\nCOMPARISON_OP_FOR_TYPE(cmple, caml_ba_bool, bool_, CMPLE_OP)\n\n// Build dispatch table with NULL for unsupported complex types\nstatic const binary_op_table cmple_table = {.i8 = nx_c_cmple_i8,\n .u8 = nx_c_cmple_u8,\n .i16 = nx_c_cmple_i16,\n .u16 = nx_c_cmple_u16,\n .i32 = nx_c_cmple_i32,\n .i64 = nx_c_cmple_i64,\n .u32 = nx_c_cmple_u32,\n .u64 = nx_c_cmple_u64,\n .inat = nx_c_cmple_inat,\n .f16 = nx_c_cmple_f16,\n .f32 = nx_c_cmple_f32,\n .f64 = nx_c_cmple_f64,\n .c32 = NULL,\n .c64 = NULL,\n .bf16 = nx_c_cmple_bf16,\n .bool_ = nx_c_cmple_bool_,\n .i4 = nx_c_cmple_i4,\n .u4 = nx_c_cmple_u4,\n .f8e4m3 = nx_c_cmple_f8e4m3,\n .f8e5m2 = nx_c_cmple_f8e5m2};\n\n// =========== BITWISE XOR ===========\n#define XOR_OP(x, y) ((x) ^ (y))\n\n// Bitwise operations only for integer types\nBINARY_OP_FOR_TYPE(xor, int8_t, i8, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, uint8_t, u8, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, int16_t, i16, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, uint16_t, u16, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, int32_t, i32, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, int64_t, i64, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, uint32_t, u32, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, uint64_t, u64, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, intnat, inat, XOR_OP)\n\n// Float bitwise operations not well-defined\n\nINT4_OP_IMPL(xor, 1, i4, XOR_OP)\nINT4_OP_IMPL(xor, 0, u4, XOR_OP)\nBINARY_OP_FOR_TYPE(xor, caml_ba_bool, bool_, XOR_OP)\n\n// Build dispatch table with NULL for unsupported float/complex types\nstatic const binary_op_table xor_table = {.i8 = nx_c_xor_i8,\n .u8 = nx_c_xor_u8,\n .i16 = nx_c_xor_i16,\n .u16 = nx_c_xor_u16,\n .i32 = nx_c_xor_i32,\n .i64 = nx_c_xor_i64,\n .u32 = nx_c_xor_u32,\n .u64 = nx_c_xor_u64,\n .inat = nx_c_xor_inat,\n .f16 = NULL,\n .f32 = NULL,\n .f64 = NULL,\n .c32 = NULL,\n .c64 = NULL,\n .bf16 = NULL,\n .bool_ = nx_c_xor_bool_,\n .i4 = nx_c_xor_i4,\n .u4 = nx_c_xor_u4,\n .f8e4m3 = NULL,\n .f8e5m2 = NULL};\n\n// =========== BITWISE OR ===========\n#define OR_OP(x, y) ((x) | (y))\n\nBINARY_OP_FOR_TYPE(or, int8_t, i8, OR_OP)\nBINARY_OP_FOR_TYPE(or, uint8_t, u8, OR_OP)\nBINARY_OP_FOR_TYPE(or, int16_t, i16, OR_OP)\nBINARY_OP_FOR_TYPE(or, uint16_t, u16, OR_OP)\nBINARY_OP_FOR_TYPE(or, int32_t, i32, OR_OP)\nBINARY_OP_FOR_TYPE(or, int64_t, i64, OR_OP)\nBINARY_OP_FOR_TYPE(or, uint32_t, u32, OR_OP)\nBINARY_OP_FOR_TYPE(or, uint64_t, u64, OR_OP)\nBINARY_OP_FOR_TYPE(or, intnat, inat, OR_OP)\n\n// Float bitwise operations not well-defined\n\nINT4_OP_IMPL(or, 1, i4, OR_OP)\nINT4_OP_IMPL(or, 0, u4, OR_OP)\nBINARY_OP_FOR_TYPE(or, caml_ba_bool, bool_, OR_OP)\n\n// Build dispatch table with NULL for unsupported float/complex types\nstatic const binary_op_table or_table = {.i8 = nx_c_or_i8,\n .u8 = nx_c_or_u8,\n .i16 = nx_c_or_i16,\n .u16 = nx_c_or_u16,\n .i32 = nx_c_or_i32,\n .i64 = nx_c_or_i64,\n .u32 = nx_c_or_u32,\n .u64 = nx_c_or_u64,\n .inat = nx_c_or_inat,\n .f16 = NULL,\n .f32 = NULL,\n .f64 = NULL,\n .c32 = NULL,\n .c64 = NULL,\n .bf16 = NULL,\n .bool_ = nx_c_or_bool_,\n .i4 = nx_c_or_i4,\n .u4 = nx_c_or_u4,\n .f8e4m3 = NULL,\n .f8e5m2 = NULL};\n\n// =========== BITWISE AND ===========\n#define AND_OP(x, y) ((x) & (y))\n\nBINARY_OP_FOR_TYPE(and, int8_t, i8, AND_OP)\nBINARY_OP_FOR_TYPE(and, uint8_t, u8, AND_OP)\nBINARY_OP_FOR_TYPE(and, int16_t, i16, AND_OP)\nBINARY_OP_FOR_TYPE(and, uint16_t, u16, AND_OP)\nBINARY_OP_FOR_TYPE(and, int32_t, i32, AND_OP)\nBINARY_OP_FOR_TYPE(and, int64_t, i64, AND_OP)\nBINARY_OP_FOR_TYPE(and, uint32_t, u32, AND_OP)\nBINARY_OP_FOR_TYPE(and, uint64_t, u64, AND_OP)\nBINARY_OP_FOR_TYPE(and, intnat, inat, AND_OP)\n\n// Float bitwise operations not well-defined\n\nINT4_OP_IMPL(and, 1, i4, AND_OP)\nINT4_OP_IMPL(and, 0, u4, AND_OP)\nBINARY_OP_FOR_TYPE(and, caml_ba_bool, bool_, AND_OP)\n\n// Build dispatch table with NULL for unsupported float/complex types\nstatic const binary_op_table and_table = {.i8 = nx_c_and_i8,\n .u8 = nx_c_and_u8,\n .i16 = nx_c_and_i16,\n .u16 = nx_c_and_u16,\n .i32 = nx_c_and_i32,\n .i64 = nx_c_and_i64,\n .u32 = nx_c_and_u32,\n .u64 = nx_c_and_u64,\n .inat = nx_c_and_inat,\n .f16 = NULL,\n .f32 = NULL,\n .f64 = NULL,\n .c32 = NULL,\n .c64 = NULL,\n .bf16 = NULL,\n .bool_ = nx_c_and_bool_,\n .i4 = nx_c_and_i4,\n .u4 = nx_c_and_u4,\n .f8e4m3 = NULL,\n .f8e5m2 = NULL};\n\n// Shared dispatch infrastructure\n\n// Generic dispatch function for binary operations\nstatic void dispatch_binary_op(value v_x, value v_y, value v_z,\n const binary_op_table *table,\n const char *op_name) {\n // Extract ndarrays from FFI tensors\n ndarray_t x = extract_ndarray(v_x);\n ndarray_t y = extract_ndarray(v_y);\n ndarray_t z = extract_ndarray(v_z);\n\n // Check shapes match\n if (x.ndim != y.ndim || x.ndim != z.ndim) {\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"shape mismatch\");\n }\n for (int i = 0; i < x.ndim; i++) {\n if (x.shape[i] != y.shape[i] || x.shape[i] != z.shape[i]) {\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"shape mismatch\");\n }\n }\n\n // Get bigarray kind from the data field\n value v_x_data = Field(v_x, FFI_TENSOR_DATA);\n value v_y_data = Field(v_y, FFI_TENSOR_DATA);\n value v_z_data = Field(v_z, FFI_TENSOR_DATA);\n\n struct caml_ba_array *ba = Caml_ba_array_val(v_x_data);\n int kind = nx_buffer_get_kind(ba);\n\n // Check kinds match for y and z\n int kind_y = nx_buffer_get_kind(Caml_ba_array_val(v_y_data));\n int kind_z = nx_buffer_get_kind(Caml_ba_array_val(v_z_data));\n if (kind != kind_y || kind != kind_z) {\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"dtype mismatch\");\n }\n\n // Select operation based on dtype\n binary_op_t op = NULL;\n switch (kind) {\n case CAML_BA_SINT8:\n op = table->i8;\n break;\n case CAML_BA_UINT8:\n op = table->u8;\n break;\n case CAML_BA_SINT16:\n op = table->i16;\n break;\n case CAML_BA_UINT16:\n op = table->u16;\n break;\n case CAML_BA_INT32:\n op = table->i32;\n break;\n case CAML_BA_INT64:\n op = table->i64;\n break;\n case NX_BA_UINT32:\n op = table->u32;\n break;\n case NX_BA_UINT64:\n op = table->u64;\n break;\n case CAML_BA_CAML_INT:\n case CAML_BA_NATIVE_INT:\n op = table->inat;\n break;\n case CAML_BA_FLOAT16:\n op = table->f16;\n break;\n case CAML_BA_FLOAT32:\n op = table->f32;\n break;\n case CAML_BA_FLOAT64:\n op = table->f64;\n break;\n case CAML_BA_COMPLEX32:\n op = table->c32;\n break;\n case CAML_BA_COMPLEX64:\n op = table->c64;\n break;\n case NX_BA_BFLOAT16:\n op = table->bf16;\n break;\n case NX_BA_BOOL:\n op = table->bool_;\n break;\n case NX_BA_INT4:\n op = table->i4;\n break;\n case NX_BA_UINT4:\n op = table->u4;\n break;\n case NX_BA_FP8_E4M3:\n op = table->f8e4m3;\n break;\n case NX_BA_FP8_E5M2:\n op = table->f8e5m2;\n break;\n default:\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"dispatch_binary_op: unsupported dtype\");\n }\n\n if (!op) {\n char msg[256];\n snprintf(msg, sizeof(msg), \"%s: operation not supported for dtype\",\n op_name);\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(msg);\n }\n\n // Enter blocking section for potentially long computation\n caml_enter_blocking_section();\n op(&x, &y, &z);\n caml_leave_blocking_section();\n\n // Clean up if heap allocated\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n}\n\n// Generic dispatch function for comparison operations (output is always bool)\nstatic void dispatch_comparison_op(value v_x, value v_y, value v_z,\n const binary_op_table *table,\n const char *op_name) {\n // Extract ndarrays from FFI tensors\n ndarray_t x = extract_ndarray(v_x);\n ndarray_t y = extract_ndarray(v_y);\n ndarray_t z = extract_ndarray(v_z);\n\n // Check shapes match\n if (x.ndim != y.ndim || x.ndim != z.ndim) {\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"shape mismatch\");\n }\n for (int i = 0; i < x.ndim; i++) {\n if (x.shape[i] != y.shape[i] || x.shape[i] != z.shape[i]) {\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"shape mismatch\");\n }\n }\n\n // Get bigarray kind from the data field\n value v_x_data = Field(v_x, FFI_TENSOR_DATA);\n value v_y_data = Field(v_y, FFI_TENSOR_DATA);\n value v_z_data = Field(v_z, FFI_TENSOR_DATA);\n\n struct caml_ba_array *ba = Caml_ba_array_val(v_x_data);\n int kind = nx_buffer_get_kind(ba);\n\n // Check input kinds match\n int kind_y = nx_buffer_get_kind(Caml_ba_array_val(v_y_data));\n if (kind != kind_y) {\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"dtype mismatch: comparison inputs must have same dtype\");\n }\n \n // Check output is uint8\n int kind_z = nx_buffer_get_kind(Caml_ba_array_val(v_z_data));\n if (kind_z != NX_BA_BOOL) {\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"dtype mismatch: comparison output must be bool\");\n }\n\n // Select operation based on input dtype\n binary_op_t op = NULL;\n switch (kind) {\n case CAML_BA_SINT8:\n op = table->i8;\n break;\n case CAML_BA_UINT8:\n op = table->u8;\n break;\n case CAML_BA_SINT16:\n op = table->i16;\n break;\n case CAML_BA_UINT16:\n op = table->u16;\n break;\n case CAML_BA_INT32:\n op = table->i32;\n break;\n case CAML_BA_INT64:\n op = table->i64;\n break;\n case NX_BA_UINT32:\n op = table->u32;\n break;\n case NX_BA_UINT64:\n op = table->u64;\n break;\n case CAML_BA_CAML_INT:\n case CAML_BA_NATIVE_INT:\n op = table->inat;\n break;\n case CAML_BA_FLOAT16:\n op = table->f16;\n break;\n case CAML_BA_FLOAT32:\n op = table->f32;\n break;\n case CAML_BA_FLOAT64:\n op = table->f64;\n break;\n case CAML_BA_COMPLEX32:\n op = table->c32;\n break;\n case CAML_BA_COMPLEX64:\n op = table->c64;\n break;\n case NX_BA_BFLOAT16:\n op = table->bf16;\n break;\n case NX_BA_BOOL:\n op = table->bool_;\n break;\n case NX_BA_INT4:\n op = table->i4;\n break;\n case NX_BA_UINT4:\n op = table->u4;\n break;\n case NX_BA_FP8_E4M3:\n op = table->f8e4m3;\n break;\n case NX_BA_FP8_E5M2:\n op = table->f8e5m2;\n break;\n default:\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(\"dispatch_comparison_op: unsupported dtype\");\n }\n\n if (!op) {\n char msg[256];\n snprintf(msg, sizeof(msg), \"%s: operation not supported for dtype\",\n op_name);\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n caml_failwith(msg);\n }\n\n // Enter blocking section for potentially long computation\n caml_enter_blocking_section();\n op(&x, &y, &z);\n caml_leave_blocking_section();\n\n // Clean up if heap allocated\n cleanup_ndarray(&x);\n cleanup_ndarray(&y);\n cleanup_ndarray(&z);\n}\n\n// ============================================================================\n// OCaml FFI Stubs\n// ============================================================================\n\n// Macro to define FFI stub for each operation\n#define DEFINE_FFI_STUB(name) \\\n CAMLprim value caml_nx_##name(value v_x, value v_y, value v_z) { \\\n CAMLparam3(v_x, v_y, v_z); \\\n dispatch_binary_op(v_x, v_y, v_z, &name##_table, #name); \\\n CAMLreturn(Val_unit); \\\n }\n\n// Macro to define FFI stub for comparison operations\n#define DEFINE_CMP_FFI_STUB(name) \\\n CAMLprim value caml_nx_##name(value v_x, value v_y, value v_z) { \\\n CAMLparam3(v_x, v_y, v_z); \\\n dispatch_comparison_op(v_x, v_y, v_z, &name##_table, #name); \\\n CAMLreturn(Val_unit); \\\n }\n\nDEFINE_FFI_STUB(add)\nDEFINE_FFI_STUB(sub)\nDEFINE_FFI_STUB(mul)\nDEFINE_FFI_STUB(idiv)\nDEFINE_FFI_STUB(fdiv)\nDEFINE_FFI_STUB(max)\nDEFINE_FFI_STUB(min)\nDEFINE_FFI_STUB(mod)\nDEFINE_FFI_STUB(pow)\nDEFINE_FFI_STUB(atan2)\nDEFINE_CMP_FFI_STUB(cmpeq)\nDEFINE_CMP_FFI_STUB(cmpne)\nDEFINE_CMP_FFI_STUB(cmplt)\nDEFINE_CMP_FFI_STUB(cmple)\nDEFINE_FFI_STUB(xor)\nDEFINE_FFI_STUB(or)\nDEFINE_FFI_STUB(and)\n"
},
{
"path": "packages/nx/lib/backend_c/nx_c_cast.c",
"content": "/*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*/\n\n// Cast operations for nx C backend\n\n#include \n#include \n#include \n#include \n#include \n#include \n\n#include \"nx_c_shared.h\"\n\n// Type definitions for cast operations\ntypedef void (*cast_op_t)(const ndarray_t *, ndarray_t *);\n\n// Enum for dtypes to index the table\ntypedef enum {\n NX_DTYPE_I8 = 0,\n NX_DTYPE_U8,\n NX_DTYPE_I16,\n NX_DTYPE_U16,\n NX_DTYPE_I32,\n NX_DTYPE_I64,\n NX_DTYPE_U32,\n NX_DTYPE_U64,\n NX_DTYPE_INAT,\n NX_DTYPE_F16,\n NX_DTYPE_F32,\n NX_DTYPE_F64,\n NX_DTYPE_C32,\n NX_DTYPE_C64,\n NX_DTYPE_BF16,\n NX_DTYPE_BOOL,\n NX_DTYPE_I4,\n NX_DTYPE_U4,\n NX_DTYPE_F8E4M3,\n NX_DTYPE_F8E5M2,\n NX_NUM_DTYPES\n} nx_dtype;\n\n// Map caml_ba_kind to nx_dtype\nstatic nx_dtype kind_to_dtype(int kind) {\n switch (kind) {\n case CAML_BA_SINT8:\n return NX_DTYPE_I8;\n case CAML_BA_UINT8:\n return NX_DTYPE_U8;\n case CAML_BA_SINT16:\n return NX_DTYPE_I16;\n case CAML_BA_UINT16:\n return NX_DTYPE_U16;\n case CAML_BA_INT32:\n return NX_DTYPE_I32;\n case CAML_BA_INT64:\n return NX_DTYPE_I64;\n case NX_BA_UINT32:\n return NX_DTYPE_U32;\n case NX_BA_UINT64:\n return NX_DTYPE_U64;\n case CAML_BA_CAML_INT:\n case CAML_BA_NATIVE_INT:\n return NX_DTYPE_INAT;\n case CAML_BA_FLOAT16:\n return NX_DTYPE_F16;\n case CAML_BA_FLOAT32:\n return NX_DTYPE_F32;\n case CAML_BA_FLOAT64:\n return NX_DTYPE_F64;\n case CAML_BA_COMPLEX32:\n return NX_DTYPE_C32;\n case CAML_BA_COMPLEX64:\n return NX_DTYPE_C64;\n case NX_BA_BFLOAT16:\n return NX_DTYPE_BF16;\n case NX_BA_BOOL:\n return NX_DTYPE_BOOL;\n case NX_BA_INT4:\n return NX_DTYPE_I4;\n case NX_BA_UINT4:\n return NX_DTYPE_U4;\n case NX_BA_FP8_E4M3:\n return NX_DTYPE_F8E4M3;\n case NX_BA_FP8_E5M2:\n return NX_DTYPE_F8E5M2;\n default:\n return NX_NUM_DTYPES;\n }\n}\n\n// Helper to iterate over inner dimensions for unary (cast) operations\ntypedef void (*kernel_fn)(void *, void *, long, long);\n\nstatic inline void iterate_inner_dims2(const ndarray_t *x, const ndarray_t *z,\n long outer_idx, kernel_fn kernel,\n void *x_data, void *z_data) {\n if (x->ndim <= 1) {\n kernel(x_data, z_data, x->offset + outer_idx * x->strides[0],\n z->offset + outer_idx * z->strides[0]);\n return;\n }\n\n long x_base = x->offset + outer_idx * x->strides[0];\n long z_base = z->offset + outer_idx * z->strides[0];\n\n int inner_ndim = x->ndim - 1;\n int *coords = (int *)calloc(inner_ndim, sizeof(int));\n if (!coords) {\n fprintf(stderr, \"nx: iterate_inner_dims2: allocation failed\\n\");\n abort();\n }\n\n bool done = false;\n while (!done) {\n long x_off = x_base;\n long z_off = z_base;\n\n for (int i = 0; i < inner_ndim; i++) {\n x_off += coords[i] * x->strides[i + 1];\n z_off += coords[i] * z->strides[i + 1];\n }\n\n kernel(x_data, z_data, x_off, z_off);\n\n done = true;\n for (int i = inner_ndim - 1; i >= 0; i--) {\n coords[i]++;\n if (coords[i] < x->shape[i + 1]) {\n done = false;\n break;\n }\n coords[i] = 0;\n }\n }\n\n free(coords);\n}\n\n// Generic cast implementation macro\n#define CAST_IMPL(src_suffix, dst_suffix) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix(const ndarray_t *src, \\\n ndarray_t *dst) { \\\n if (!src || !dst) { \\\n fprintf(stderr, \"nx: nx_c_cast_\" #src_suffix \"_to_\" #dst_suffix \\\n \": null pointer\\n\"); \\\n abort(); \\\n } \\\n long total = total_elements_safe(src); \\\n if (total == 0) return; \\\n \\\n if (is_contiguous(src) && is_contiguous(dst)) { \\\n _Pragma(\"omp parallel for simd if(total > 1000)\") for (long i = 0; \\\n i < total; i++) { \\\n nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n src->data, dst->data, src->offset + i, dst->offset + i); \\\n } \\\n } else if (src->shape[0] > 1 && total / src->shape[0] > 50) { \\\n _Pragma(\"omp parallel for if(src->shape[0] > 4)\") for (long i = 0; \\\n i < \\\n src->shape[0]; \\\n i++) { \\\n iterate_inner_dims2(src, dst, i, \\\n nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel, \\\n src->data, dst->data); \\\n } \\\n } else { \\\n nd_copy_iterator_t it; \\\n nd_copy_iterator_init(&it, src, dst); \\\n do { \\\n long src_off, dst_off; \\\n nd_copy_iterator_get_offsets(&it, &src_off, &dst_off); \\\n nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n src->data, dst->data, src->offset + src_off, \\\n dst->offset + dst_off); \\\n } while (nd_copy_iterator_next(&it)); \\\n nd_copy_iterator_destroy(&it); \\\n } \\\n }\n\n// Note: Assumes is_fully_contiguous accepts 3 arguments, with NULL for y.\n// Assume shared.h has overload or ignore NULL.\n\n// Standard real types (integer and float, excluding bool, low prec, complex,\n// packed)\n#define STANDARD_REAL_TYPES \\\n SR_TYPE(i8, int8_t) \\\n SR_TYPE(u8, uint8_t) \\\n SR_TYPE(i16, int16_t) \\\n SR_TYPE(u16, uint16_t) \\\n SR_TYPE(i32, int32_t) \\\n SR_TYPE(i64, int64_t) \\\n SR_TYPE(u32, uint32_t) \\\n SR_TYPE(u64, uint64_t) \\\n SR_TYPE(inat, intnat) \\\n SR_TYPE(f32, float) \\\n SR_TYPE(f64, double)\n\n// Low precision float types\n#define LOW_PREC_TYPES \\\n LP_TYPE(f16, uint16_t, half_to_float, float_to_half) \\\n LP_TYPE(bf16, caml_ba_bfloat16, bfloat16_to_float, float_to_bfloat16) \\\n LP_TYPE(f8e4m3, caml_ba_fp8_e4m3, fp8_e4m3_to_float, float_to_fp8_e4m3) \\\n LP_TYPE(f8e5m2, caml_ba_fp8_e5m2, fp8_e5m2_to_float, float_to_fp8_e5m2)\n\n// Complex types\n#define COMPLEX_TYPES \\\n CP_TYPE(c32, complex32, crealf(src[src_off]), cimagf(src[src_off]), float) \\\n CP_TYPE(c64, complex64, creal(src[src_off]), cimag(src[src_off]), double)\n\n// Generate cast for standard real to standard real\n\n#define GEN_CAST_STANDARD_TO_STANDARD(src_suffix, dst_suffix, src_t, dst_t) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n dst[dst_off] = (dst_t)src[src_off]; \\\n } \\\n CAST_IMPL(src_suffix, dst_suffix)\n\n// Generate all standard-to-standard cast combinations\n#define GEN_ALL_STANDARD_CASTS(src_suffix, src_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, i8, src_t, int8_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, u8, src_t, uint8_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, i16, src_t, int16_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, u16, src_t, uint16_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, i32, src_t, int32_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, i64, src_t, int64_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, u32, src_t, uint32_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, u64, src_t, uint64_t) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, inat, src_t, intnat) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, f32, src_t, float) \\\n GEN_CAST_STANDARD_TO_STANDARD(src_suffix, f64, src_t, double)\n\n// Generate casts for each standard type\nGEN_ALL_STANDARD_CASTS(i8, int8_t)\nGEN_ALL_STANDARD_CASTS(u8, uint8_t)\nGEN_ALL_STANDARD_CASTS(i16, int16_t)\nGEN_ALL_STANDARD_CASTS(u16, uint16_t)\nGEN_ALL_STANDARD_CASTS(i32, int32_t)\nGEN_ALL_STANDARD_CASTS(i64, int64_t)\nGEN_ALL_STANDARD_CASTS(u32, uint32_t)\nGEN_ALL_STANDARD_CASTS(u64, uint64_t)\nGEN_ALL_STANDARD_CASTS(inat, intnat)\nGEN_ALL_STANDARD_CASTS(f32, float)\nGEN_ALL_STANDARD_CASTS(f64, double)\n\n// Generate cast for bool to standard real\n\n#define GEN_CAST_BOOL_TO_STANDARD(dst_suffix, dst_t) \\\n static void nx_c_cast_bool_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n dst[dst_off] = (dst_t)src[src_off]; \\\n } \\\n CAST_IMPL(bool, dst_suffix)\n\n#define SR_TYPE(dst_suffix, dst_t) GEN_CAST_BOOL_TO_STANDARD(dst_suffix, dst_t)\nSTANDARD_REAL_TYPES\n#undef SR_TYPE\n\n// Generate cast for standard real to bool\n\n#define GEN_CAST_STANDARD_TO_BOOL(src_suffix, src_t) \\\n static void nx_c_cast_##src_suffix##_to_bool_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n uint8_t *dst = (uint8_t *)dst_data; \\\n dst[dst_off] = (src[src_off] != 0) ? 1 : 0; \\\n } \\\n CAST_IMPL(src_suffix, bool)\n\n#define SR_TYPE(src_suffix, src_t) GEN_CAST_STANDARD_TO_BOOL(src_suffix, src_t)\nSTANDARD_REAL_TYPES\n#undef SR_TYPE\n\n// Bool to bool\n\nstatic void nx_c_cast_bool_to_bool_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n uint8_t *src = (uint8_t *)src_data;\n uint8_t *dst = (uint8_t *)dst_data;\n dst[dst_off] = (src[src_off] != 0) ? 1 : 0;\n}\nCAST_IMPL(bool, bool)\n\n// Generate cast for low prec to standard real\n\n#define GEN_CAST_LP_TO_STANDARD(src_suffix, dst_suffix, src_t, dst_t, \\\n TO_FLOAT) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n float temp = TO_FLOAT(src[src_off]); \\\n dst[dst_off] = (dst_t)temp; \\\n } \\\n CAST_IMPL(src_suffix, dst_suffix)\n\n// Generate all low-precision to standard casts\n#define GEN_ALL_LP_TO_STANDARD_CASTS(src_suffix, src_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, i8, src_t, int8_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, u8, src_t, uint8_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, i16, src_t, int16_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, u16, src_t, uint16_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, i32, src_t, int32_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, i64, src_t, int64_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, u32, src_t, uint32_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, u64, src_t, uint64_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, inat, src_t, intnat, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, f32, src_t, float, TO_FLOAT) \\\n GEN_CAST_LP_TO_STANDARD(src_suffix, f64, src_t, double, TO_FLOAT)\n\n// Generate casts for each low-precision type\nGEN_ALL_LP_TO_STANDARD_CASTS(f16, uint16_t, half_to_float)\nGEN_ALL_LP_TO_STANDARD_CASTS(bf16, caml_ba_bfloat16, bfloat16_to_float)\nGEN_ALL_LP_TO_STANDARD_CASTS(f8e4m3, caml_ba_fp8_e4m3, fp8_e4m3_to_float)\nGEN_ALL_LP_TO_STANDARD_CASTS(f8e5m2, caml_ba_fp8_e5m2, fp8_e5m2_to_float)\n\n// Generate cast for standard real to low prec\n\n#define GEN_CAST_STANDARD_TO_LP(src_suffix, dst_suffix, src_t, dst_t, \\\n FROM_FLOAT) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n float temp = (float)src[src_off]; \\\n dst[dst_off] = FROM_FLOAT(temp); \\\n } \\\n CAST_IMPL(src_suffix, dst_suffix)\n\n// Generate all standard to low-precision casts\n#define GEN_ALL_STANDARD_TO_LP_CASTS(dst_suffix, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(i8, dst_suffix, int8_t, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(u8, dst_suffix, uint8_t, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(i16, dst_suffix, int16_t, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(u16, dst_suffix, uint16_t, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(i32, dst_suffix, int32_t, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(i64, dst_suffix, int64_t, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(u32, dst_suffix, uint32_t, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(u64, dst_suffix, uint64_t, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(inat, dst_suffix, intnat, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(f32, dst_suffix, float, dst_t, FROM_FLOAT) \\\n GEN_CAST_STANDARD_TO_LP(f64, dst_suffix, double, dst_t, FROM_FLOAT)\n\n// Generate casts for each low-precision type\nGEN_ALL_STANDARD_TO_LP_CASTS(f16, uint16_t, float_to_half)\nGEN_ALL_STANDARD_TO_LP_CASTS(bf16, caml_ba_bfloat16, float_to_bfloat16)\nGEN_ALL_STANDARD_TO_LP_CASTS(f8e4m3, caml_ba_fp8_e4m3, float_to_fp8_e4m3)\nGEN_ALL_STANDARD_TO_LP_CASTS(f8e5m2, caml_ba_fp8_e5m2, float_to_fp8_e5m2)\n\n// Bool to low prec\n#define GEN_CAST_BOOL_TO_LP(dst_suffix, dst_t, FROM_FLOAT) \\\n static void nx_c_cast_bool_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n float temp = (src[src_off] != 0) ? 1.0f : 0.0f; \\\n dst[dst_off] = FROM_FLOAT(temp); \\\n } \\\n CAST_IMPL(bool, dst_suffix)\n\nGEN_CAST_BOOL_TO_LP(f16, uint16_t, float_to_half)\nGEN_CAST_BOOL_TO_LP(bf16, caml_ba_bfloat16, float_to_bfloat16)\nGEN_CAST_BOOL_TO_LP(f8e4m3, caml_ba_fp8_e4m3, float_to_fp8_e4m3)\nGEN_CAST_BOOL_TO_LP(f8e5m2, caml_ba_fp8_e5m2, float_to_fp8_e5m2)\n\n// Generate cast for low prec to bool\n\n#define GEN_CAST_LP_TO_BOOL(src_suffix, src_t, TO_FLOAT) \\\n static void nx_c_cast_##src_suffix##_to_bool_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n uint8_t *dst = (uint8_t *)dst_data; \\\n float temp = TO_FLOAT(src[src_off]); \\\n dst[dst_off] = (temp != 0.0f) ? 1 : 0; \\\n } \\\n CAST_IMPL(src_suffix, bool)\n\n#define LP_TYPE(suffix, t, to_f, from_f) GEN_CAST_LP_TO_BOOL(suffix, t, to_f)\nLOW_PREC_TYPES\n#undef LP_TYPE\n\n// Generate cast for low prec to low prec\n\n#define GEN_CAST_LP_TO_LP(src_suffix, dst_suffix, src_t, dst_t, TO_FLOAT, \\\n FROM_FLOAT) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n float temp = TO_FLOAT(src[src_off]); \\\n dst[dst_off] = FROM_FLOAT(temp); \\\n } \\\n CAST_IMPL(src_suffix, dst_suffix)\n\n// Generate all low-precision to low-precision casts\n// Identity casts for low-precision types\nstatic void nx_c_cast_f16_to_f16_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n uint16_t *src = (uint16_t *)src_data;\n uint16_t *dst = (uint16_t *)dst_data;\n dst[dst_off] = src[src_off];\n}\nCAST_IMPL(f16, f16)\n\nGEN_CAST_LP_TO_LP(f16, bf16, uint16_t, caml_ba_bfloat16, half_to_float,\n float_to_bfloat16)\nGEN_CAST_LP_TO_LP(f16, f8e4m3, uint16_t, caml_ba_fp8_e4m3, half_to_float,\n float_to_fp8_e4m3)\nGEN_CAST_LP_TO_LP(f16, f8e5m2, uint16_t, caml_ba_fp8_e5m2, half_to_float,\n float_to_fp8_e5m2)\nGEN_CAST_LP_TO_LP(bf16, f16, caml_ba_bfloat16, uint16_t, bfloat16_to_float,\n float_to_half)\nGEN_CAST_LP_TO_LP(bf16, f8e4m3, caml_ba_bfloat16, caml_ba_fp8_e4m3,\n bfloat16_to_float, float_to_fp8_e4m3)\nGEN_CAST_LP_TO_LP(bf16, f8e5m2, caml_ba_bfloat16, caml_ba_fp8_e5m2,\n bfloat16_to_float, float_to_fp8_e5m2)\nGEN_CAST_LP_TO_LP(f8e4m3, f16, caml_ba_fp8_e4m3, uint16_t, fp8_e4m3_to_float,\n float_to_half)\nGEN_CAST_LP_TO_LP(f8e4m3, bf16, caml_ba_fp8_e4m3, caml_ba_bfloat16,\n fp8_e4m3_to_float, float_to_bfloat16)\nGEN_CAST_LP_TO_LP(f8e4m3, f8e5m2, caml_ba_fp8_e4m3, caml_ba_fp8_e5m2,\n fp8_e4m3_to_float, float_to_fp8_e5m2)\nGEN_CAST_LP_TO_LP(f8e5m2, f16, caml_ba_fp8_e5m2, uint16_t, fp8_e5m2_to_float,\n float_to_half)\nGEN_CAST_LP_TO_LP(f8e5m2, bf16, caml_ba_fp8_e5m2, caml_ba_bfloat16,\n fp8_e5m2_to_float, float_to_bfloat16)\nGEN_CAST_LP_TO_LP(f8e5m2, f8e4m3, caml_ba_fp8_e5m2, caml_ba_fp8_e4m3,\n fp8_e5m2_to_float, float_to_fp8_e4m3)\n\n// Generate cast for complex to standard real\n\n#define GEN_CAST_CP_TO_STANDARD(src_suffix, dst_suffix, src_t, dst_t, RE_FN) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n double temp = RE_FN; \\\n dst[dst_off] = (dst_t)temp; \\\n } \\\n CAST_IMPL(src_suffix, dst_suffix)\n\n// Generate all complex to standard casts\n#define GEN_ALL_CP_TO_STANDARD_CASTS(src_suffix, src_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, i8, src_t, int8_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, u8, src_t, uint8_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, i16, src_t, int16_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, u16, src_t, uint16_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, i32, src_t, int32_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, i64, src_t, int64_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, u32, src_t, uint32_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, u64, src_t, uint64_t, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, inat, src_t, intnat, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, f32, src_t, float, RE_FN) \\\n GEN_CAST_CP_TO_STANDARD(src_suffix, f64, src_t, double, RE_FN)\n\nGEN_ALL_CP_TO_STANDARD_CASTS(c32, complex32, crealf(src[src_off]))\nGEN_ALL_CP_TO_STANDARD_CASTS(c64, complex64, creal(src[src_off]))\n\n// Generate cast for complex to bool\n\n#define GEN_CAST_CP_TO_BOOL(src_suffix, src_t, RE_FN, IM_FN) \\\n static void nx_c_cast_##src_suffix##_to_bool_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n uint8_t *dst = (uint8_t *)dst_data; \\\n double re = RE_FN; \\\n double im = IM_FN; \\\n dst[dst_off] = (re != 0 || im != 0) ? 1 : 0; \\\n } \\\n CAST_IMPL(src_suffix, bool)\n\n#define CP_TYPE(suffix, t, re_fn, im_fn, base_t) \\\n GEN_CAST_CP_TO_BOOL(suffix, t, re_fn, im_fn)\nCOMPLEX_TYPES\n#undef CP_TYPE\n\n// Generate cast for complex to low prec float\n\n#define GEN_CAST_CP_TO_LP(src_suffix, dst_suffix, src_t, dst_t, RE_FN, \\\n FROM_FLOAT) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n float temp = RE_FN; \\\n dst[dst_off] = FROM_FLOAT(temp); \\\n } \\\n CAST_IMPL(src_suffix, dst_suffix)\n\n// Generate all complex to low-precision casts\n#define GEN_ALL_CP_TO_LP_CASTS(src_suffix, src_t, RE_FN) \\\n GEN_CAST_CP_TO_LP(src_suffix, f16, src_t, uint16_t, RE_FN, float_to_half) \\\n GEN_CAST_CP_TO_LP(src_suffix, bf16, src_t, caml_ba_bfloat16, RE_FN, \\\n float_to_bfloat16) \\\n GEN_CAST_CP_TO_LP(src_suffix, f8e4m3, src_t, caml_ba_fp8_e4m3, RE_FN, \\\n float_to_fp8_e4m3) \\\n GEN_CAST_CP_TO_LP(src_suffix, f8e5m2, src_t, caml_ba_fp8_e5m2, RE_FN, \\\n float_to_fp8_e5m2)\n\nGEN_ALL_CP_TO_LP_CASTS(c32, complex32, crealf(src[src_off]))\nGEN_ALL_CP_TO_LP_CASTS(c64, complex64, creal(src[src_off]))\n\n// Generate cast for standard real to c32/c64\n\n#define GEN_CAST_STANDARD_TO_C32_C64(src_suffix, dst_suffix, src_t, dst_t, \\\n BASE_T) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n BASE_T temp = (BASE_T)src[src_off]; \\\n dst[dst_off] = temp + 0.0 * I; \\\n } \\\n CAST_IMPL(src_suffix, dst_suffix)\n\n\n// Generate all standard to complex casts for c32/c64\n#define GEN_ALL_STANDARD_TO_C32_C64_CASTS(dst_suffix, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(i8, dst_suffix, int8_t, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(u8, dst_suffix, uint8_t, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(i16, dst_suffix, int16_t, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(u16, dst_suffix, uint16_t, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(i32, dst_suffix, int32_t, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(i64, dst_suffix, int64_t, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(u32, dst_suffix, uint32_t, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(u64, dst_suffix, uint64_t, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(inat, dst_suffix, intnat, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(f32, dst_suffix, float, dst_t, BASE_T) \\\n GEN_CAST_STANDARD_TO_C32_C64(f64, dst_suffix, double, dst_t, BASE_T)\n\nGEN_ALL_STANDARD_TO_C32_C64_CASTS(c32, complex32, float)\nGEN_ALL_STANDARD_TO_C32_C64_CASTS(c64, complex64, double)\n\n\n// Removed - Already generated above\n\n// Bool to c32/c64\n\n#define GEN_CAST_BOOL_TO_CP(dst_suffix, dst_t, BASE_T) \\\n static void nx_c_cast_bool_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n BASE_T temp = (BASE_T)src[src_off]; \\\n dst[dst_off] = temp + 0.0 * I; \\\n } \\\n CAST_IMPL(bool, dst_suffix)\n\nGEN_CAST_BOOL_TO_CP(c32, complex32, float)\nGEN_CAST_BOOL_TO_CP(c64, complex64, double)\n\n// Low prec to c32/c64\n\n#define GEN_CAST_LP_TO_CP(src_suffix, src_t, dst_suffix, dst_t, BASE_T, \\\n TO_FLOAT) \\\n static void nx_c_cast_##src_suffix##_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n BASE_T temp = (BASE_T)TO_FLOAT(src[src_off]); \\\n dst[dst_off] = temp + 0.0 * I; \\\n } \\\n CAST_IMPL(src_suffix, dst_suffix)\n\n#define DEFINE_CASTS_LP_TO_CP(src_suffix, src_t, TO_FLOAT) \\\n GEN_CAST_LP_TO_CP(src_suffix, src_t, c32, complex32, float, TO_FLOAT) \\\n GEN_CAST_LP_TO_CP(src_suffix, src_t, c64, complex64, double, TO_FLOAT)\n\n#define LP_TYPE(suffix, t, to_f, from_f) DEFINE_CASTS_LP_TO_CP(suffix, t, to_f)\nLOW_PREC_TYPES\n#undef LP_TYPE\n\n// Complex to complex pairs (individual)\n\nstatic void nx_c_cast_c32_to_c32_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n complex32 *src = (complex32 *)src_data;\n complex32 *dst = (complex32 *)dst_data;\n dst[dst_off] = src[src_off];\n}\nCAST_IMPL(c32, c32)\n\nstatic void nx_c_cast_c32_to_c64_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n complex32 *src = (complex32 *)src_data;\n complex64 *dst = (complex64 *)dst_data;\n dst[dst_off] = (complex64)src[src_off];\n}\nCAST_IMPL(c32, c64)\n\n\nstatic void nx_c_cast_c64_to_c32_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n complex64 *src = (complex64 *)src_data;\n complex32 *dst = (complex32 *)dst_data;\n dst[dst_off] = (complex32)src[src_off];\n}\nCAST_IMPL(c64, c32)\n\nstatic void nx_c_cast_c64_to_c64_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n complex64 *src = (complex64 *)src_data;\n complex64 *dst = (complex64 *)dst_data;\n dst[dst_off] = src[src_off];\n}\nCAST_IMPL(c64, c64)\n\n\n// Generate cast for i4 to standard real\n\n#define GEN_CAST_I4_TO_STANDARD(dst_suffix, dst_t) \\\n static void nx_c_cast_i4_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n long byte_off = src_off / 2; \\\n int nib_off = src_off % 2; \\\n int a = nib_off ? ((int8_t)src[byte_off] >> 4) \\\n : (int8_t)((src[byte_off] & 0x0F) << 4) >> 4; \\\n dst[dst_off] = (dst_t)a; \\\n } \\\n CAST_IMPL(i4, dst_suffix)\n\n#define SR_TYPE(dst_suffix, dst_t) GEN_CAST_I4_TO_STANDARD(dst_suffix, dst_t)\nSTANDARD_REAL_TYPES\n#undef SR_TYPE\n\n// Generate cast for u4 to standard real\n\n#define GEN_CAST_U4_TO_STANDARD(dst_suffix, dst_t) \\\n static void nx_c_cast_u4_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n long byte_off = src_off / 2; \\\n int nib_off = src_off % 2; \\\n int a = nib_off ? (src[byte_off] >> 4) & 0x0F : src[byte_off] & 0x0F; \\\n dst[dst_off] = (dst_t)a; \\\n } \\\n CAST_IMPL(u4, dst_suffix)\n\n#define SR_TYPE(dst_suffix, dst_t) GEN_CAST_U4_TO_STANDARD(dst_suffix, dst_t)\nSTANDARD_REAL_TYPES\n#undef SR_TYPE\n\n// i4 to bool\n\nstatic void nx_c_cast_i4_to_bool_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n uint8_t *src = (uint8_t *)src_data;\n uint8_t *dst = (uint8_t *)dst_data;\n long byte_off = src_off / 2;\n int nib_off = src_off % 2;\n int a = nib_off ? ((int8_t)src[byte_off] >> 4)\n : (int8_t)((src[byte_off] & 0x0F) << 4) >> 4;\n dst[dst_off] = (a != 0) ? 1 : 0;\n}\nCAST_IMPL(i4, bool)\n\n// u4 to bool\n\nstatic void nx_c_cast_u4_to_bool_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n uint8_t *src = (uint8_t *)src_data;\n uint8_t *dst = (uint8_t *)dst_data;\n long byte_off = src_off / 2;\n int nib_off = src_off % 2;\n int a = nib_off ? (src[byte_off] >> 4) & 0x0F : src[byte_off] & 0x0F;\n dst[dst_off] = (a != 0) ? 1 : 0;\n}\nCAST_IMPL(u4, bool)\n\n// Generate cast for i4 to low prec\n\n#define GEN_CAST_I4_TO_LP(dst_suffix, dst_t, FROM_FLOAT) \\\n static void nx_c_cast_i4_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n long byte_off = src_off / 2; \\\n int nib_off = src_off % 2; \\\n int a = nib_off ? ((int8_t)src[byte_off] >> 4) \\\n : (int8_t)((src[byte_off] & 0x0F) << 4) >> 4; \\\n float temp = (float)a; \\\n dst[dst_off] = FROM_FLOAT(temp); \\\n } \\\n CAST_IMPL(i4, dst_suffix)\n\n#define LP_TYPE(suffix, t, to_f, from_f) GEN_CAST_I4_TO_LP(suffix, t, from_f)\nLOW_PREC_TYPES\n#undef LP_TYPE\n\n// Generate cast for u4 to low prec\n\n#define GEN_CAST_U4_TO_LP(dst_suffix, dst_t, FROM_FLOAT) \\\n static void nx_c_cast_u4_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n long byte_off = src_off / 2; \\\n int nib_off = src_off % 2; \\\n int a = nib_off ? (src[byte_off] >> 4) & 0x0F : src[byte_off] & 0x0F; \\\n float temp = (float)a; \\\n dst[dst_off] = FROM_FLOAT(temp); \\\n } \\\n CAST_IMPL(u4, dst_suffix)\n\n#define LP_TYPE(suffix, t, to_f, from_f) GEN_CAST_U4_TO_LP(suffix, t, from_f)\nLOW_PREC_TYPES\n#undef LP_TYPE\n\n// Generate cast for i4 to c32/c64\n\n#define GEN_CAST_I4_TO_CP(dst_suffix, dst_t, BASE_T) \\\n static void nx_c_cast_i4_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n long byte_off = src_off / 2; \\\n int nib_off = src_off % 2; \\\n int a = nib_off ? ((int8_t)src[byte_off] >> 4) \\\n : (int8_t)((src[byte_off] & 0x0F) << 4) >> 4; \\\n BASE_T temp = (BASE_T)a; \\\n dst[dst_off] = temp + 0.0 * I; \\\n } \\\n CAST_IMPL(i4, dst_suffix)\n\nGEN_CAST_I4_TO_CP(c32, complex32, float)\nGEN_CAST_I4_TO_CP(c64, complex64, double)\n\n// Similar for u4 to c32/c64\n\n#define GEN_CAST_U4_TO_CP(dst_suffix, dst_t, BASE_T) \\\n static void nx_c_cast_u4_to_##dst_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n dst_t *dst = (dst_t *)dst_data; \\\n long byte_off = src_off / 2; \\\n int nib_off = src_off % 2; \\\n int a = nib_off ? (src[byte_off] >> 4) & 0x0F : src[byte_off] & 0x0F; \\\n BASE_T temp = (BASE_T)a; \\\n dst[dst_off] = temp + 0.0 * I; \\\n } \\\n CAST_IMPL(u4, dst_suffix)\n\nGEN_CAST_U4_TO_CP(c32, complex32, float)\nGEN_CAST_U4_TO_CP(c64, complex64, double)\n\n// Generate cast for standard real to i4/u4\n\n#define GEN_CAST_STANDARD_TO_PACKED(src_suffix, src_t, packed_suffix) \\\n static void nx_c_cast_##src_suffix##_to_##packed_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n uint8_t *dst = (uint8_t *)dst_data; \\\n long byte_off = dst_off / 2; \\\n int nib_off = dst_off % 2; \\\n int res = (int)src[src_off]; \\\n uint8_t nib = (uint8_t)res & 0x0F; \\\n if (nib_off) { \\\n dst[byte_off] = (dst[byte_off] & 0x0F) | (nib << 4); \\\n } else { \\\n dst[byte_off] = (dst[byte_off] & 0xF0) | nib; \\\n } \\\n } \\\n CAST_IMPL(src_suffix, packed_suffix)\n\n// Generate standard to packed casts\nGEN_CAST_STANDARD_TO_PACKED(i8, int8_t, i4)\nGEN_CAST_STANDARD_TO_PACKED(u8, uint8_t, i4)\nGEN_CAST_STANDARD_TO_PACKED(i16, int16_t, i4)\nGEN_CAST_STANDARD_TO_PACKED(u16, uint16_t, i4)\nGEN_CAST_STANDARD_TO_PACKED(i32, int32_t, i4)\nGEN_CAST_STANDARD_TO_PACKED(i64, int64_t, i4)\nGEN_CAST_STANDARD_TO_PACKED(u32, uint32_t, i4)\nGEN_CAST_STANDARD_TO_PACKED(u64, uint64_t, i4)\nGEN_CAST_STANDARD_TO_PACKED(inat, intnat, i4)\nGEN_CAST_STANDARD_TO_PACKED(f32, float, i4)\nGEN_CAST_STANDARD_TO_PACKED(f64, double, i4)\n\nGEN_CAST_STANDARD_TO_PACKED(i8, int8_t, u4)\nGEN_CAST_STANDARD_TO_PACKED(u8, uint8_t, u4)\nGEN_CAST_STANDARD_TO_PACKED(i16, int16_t, u4)\nGEN_CAST_STANDARD_TO_PACKED(u16, uint16_t, u4)\nGEN_CAST_STANDARD_TO_PACKED(i32, int32_t, u4)\nGEN_CAST_STANDARD_TO_PACKED(i64, int64_t, u4)\nGEN_CAST_STANDARD_TO_PACKED(u32, uint32_t, u4)\nGEN_CAST_STANDARD_TO_PACKED(u64, uint64_t, u4)\nGEN_CAST_STANDARD_TO_PACKED(inat, intnat, u4)\nGEN_CAST_STANDARD_TO_PACKED(f32, float, u4)\nGEN_CAST_STANDARD_TO_PACKED(f64, double, u4)\n\n// Bool to i4/u4\n\n#define GEN_CAST_BOOL_TO_PACKED(packed_suffix) \\\n static void nx_c_cast_bool_to_##packed_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n uint8_t *src = (uint8_t *)src_data; \\\n uint8_t *dst = (uint8_t *)dst_data; \\\n long byte_off = dst_off / 2; \\\n int nib_off = dst_off % 2; \\\n int res = (int)src[src_off]; \\\n uint8_t nib = (uint8_t)res & 0x0F; \\\n if (nib_off) { \\\n dst[byte_off] = (dst[byte_off] & 0x0F) | (nib << 4); \\\n } else { \\\n dst[byte_off] = (dst[byte_off] & 0xF0) | nib; \\\n } \\\n } \\\n CAST_IMPL(bool, packed_suffix)\n\nGEN_CAST_BOOL_TO_PACKED(i4)\nGEN_CAST_BOOL_TO_PACKED(u4)\n\n// Low prec to i4/u4\n\n#define GEN_CAST_LP_TO_PACKED(src_suffix, src_t, TO_FLOAT, packed_suffix) \\\n static void nx_c_cast_##src_suffix##_to_##packed_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n uint8_t *dst = (uint8_t *)dst_data; \\\n long byte_off = dst_off / 2; \\\n int nib_off = dst_off % 2; \\\n float temp = TO_FLOAT(src[src_off]); \\\n int res = (int)temp; \\\n uint8_t nib = (uint8_t)res & 0x0F; \\\n if (nib_off) { \\\n dst[byte_off] = (dst[byte_off] & 0x0F) | (nib << 4); \\\n } else { \\\n dst[byte_off] = (dst[byte_off] & 0xF0) | nib; \\\n } \\\n } \\\n CAST_IMPL(src_suffix, packed_suffix)\n\n// Generate low-precision to packed casts\nGEN_CAST_LP_TO_PACKED(f16, uint16_t, half_to_float, i4)\nGEN_CAST_LP_TO_PACKED(bf16, caml_ba_bfloat16, bfloat16_to_float, i4)\nGEN_CAST_LP_TO_PACKED(f8e4m3, caml_ba_fp8_e4m3, fp8_e4m3_to_float, i4)\nGEN_CAST_LP_TO_PACKED(f8e5m2, caml_ba_fp8_e5m2, fp8_e5m2_to_float, i4)\n\nGEN_CAST_LP_TO_PACKED(f16, uint16_t, half_to_float, u4)\nGEN_CAST_LP_TO_PACKED(bf16, caml_ba_bfloat16, bfloat16_to_float, u4)\nGEN_CAST_LP_TO_PACKED(f8e4m3, caml_ba_fp8_e4m3, fp8_e4m3_to_float, u4)\nGEN_CAST_LP_TO_PACKED(f8e5m2, caml_ba_fp8_e5m2, fp8_e5m2_to_float, u4)\n\n// Complex to i4/u4\n\n#define GEN_CAST_CP_TO_PACKED(src_suffix, src_t, RE_FN, packed_suffix) \\\n static void nx_c_cast_##src_suffix##_to_##packed_suffix##_kernel( \\\n void *src_data, void *dst_data, long src_off, long dst_off) { \\\n src_t *src = (src_t *)src_data; \\\n uint8_t *dst = (uint8_t *)dst_data; \\\n long byte_off = dst_off / 2; \\\n int nib_off = dst_off % 2; \\\n float temp = RE_FN; \\\n int res = (int)temp; \\\n uint8_t nib = (uint8_t)res & 0x0F; \\\n if (nib_off) { \\\n dst[byte_off] = (dst[byte_off] & 0x0F) | (nib << 4); \\\n } else { \\\n dst[byte_off] = (dst[byte_off] & 0xF0) | nib; \\\n } \\\n } \\\n CAST_IMPL(src_suffix, packed_suffix)\n\n// Generate complex to packed casts\nGEN_CAST_CP_TO_PACKED(c32, complex32, crealf(src[src_off]), i4)\nGEN_CAST_CP_TO_PACKED(c64, complex64, creal(src[src_off]), i4)\n\nGEN_CAST_CP_TO_PACKED(c32, complex32, crealf(src[src_off]), u4)\nGEN_CAST_CP_TO_PACKED(c64, complex64, creal(src[src_off]), u4)\n\n// Identity casts for other low-precision and special types\nstatic void nx_c_cast_bf16_to_bf16_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n caml_ba_bfloat16 *src = (caml_ba_bfloat16 *)src_data;\n caml_ba_bfloat16 *dst = (caml_ba_bfloat16 *)dst_data;\n dst[dst_off] = src[src_off];\n}\nCAST_IMPL(bf16, bf16)\n\nstatic void nx_c_cast_f8e4m3_to_f8e4m3_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n caml_ba_fp8_e4m3 *src = (caml_ba_fp8_e4m3 *)src_data;\n caml_ba_fp8_e4m3 *dst = (caml_ba_fp8_e4m3 *)dst_data;\n dst[dst_off] = src[src_off];\n}\nCAST_IMPL(f8e4m3, f8e4m3)\n\nstatic void nx_c_cast_f8e5m2_to_f8e5m2_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n caml_ba_fp8_e5m2 *src = (caml_ba_fp8_e5m2 *)src_data;\n caml_ba_fp8_e5m2 *dst = (caml_ba_fp8_e5m2 *)dst_data;\n dst[dst_off] = src[src_off];\n}\nCAST_IMPL(f8e5m2, f8e5m2)\n\n// Packed to packed\n\nstatic void nx_c_cast_i4_to_i4_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n uint8_t *src = (uint8_t *)src_data;\n uint8_t *dst = (uint8_t *)dst_data;\n long byte_off = src_off / 2;\n int nib_off = src_off % 2;\n int a = nib_off ? ((int8_t)src[byte_off] >> 4)\n : (int8_t)((src[byte_off] & 0x0F) << 4) >> 4;\n uint8_t nib = (uint8_t)a & 0x0F;\n long d_byte_off = dst_off / 2;\n int d_nib_off = dst_off % 2;\n if (d_nib_off) {\n dst[d_byte_off] = (dst[d_byte_off] & 0x0F) | (nib << 4);\n } else {\n dst[d_byte_off] = (dst[d_byte_off] & 0xF0) | nib;\n }\n}\nCAST_IMPL(i4, i4)\n\nstatic void nx_c_cast_i4_to_u4_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n uint8_t *src = (uint8_t *)src_data;\n uint8_t *dst = (uint8_t *)dst_data;\n long byte_off = src_off / 2;\n int nib_off = src_off % 2;\n int a = nib_off ? ((int8_t)src[byte_off] >> 4)\n : (int8_t)((src[byte_off] & 0x0F) << 4) >> 4;\n uint8_t nib = (uint8_t)a & 0x0F;\n long d_byte_off = dst_off / 2;\n int d_nib_off = dst_off % 2;\n if (d_nib_off) {\n dst[d_byte_off] = (dst[d_byte_off] & 0x0F) | (nib << 4);\n } else {\n dst[d_byte_off] = (dst[d_byte_off] & 0xF0) | nib;\n }\n}\nCAST_IMPL(i4, u4)\n\nstatic void nx_c_cast_u4_to_i4_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n uint8_t *src = (uint8_t *)src_data;\n uint8_t *dst = (uint8_t *)dst_data;\n long byte_off = src_off / 2;\n int nib_off = src_off % 2;\n int a = nib_off ? (src[byte_off] >> 4) & 0x0F : src[byte_off] & 0x0F;\n uint8_t nib = (uint8_t)a & 0x0F;\n long d_byte_off = dst_off / 2;\n int d_nib_off = dst_off % 2;\n if (d_nib_off) {\n dst[d_byte_off] = (dst[d_byte_off] & 0x0F) | (nib << 4);\n } else {\n dst[d_byte_off] = (dst[d_byte_off] & 0xF0) | nib;\n }\n}\nCAST_IMPL(u4, i4)\n\nstatic void nx_c_cast_u4_to_u4_kernel(void *src_data, void *dst_data,\n long src_off, long dst_off) {\n uint8_t *src = (uint8_t *)src_data;\n uint8_t *dst = (uint8_t *)dst_data;\n long byte_off = src_off / 2;\n int nib_off = src_off % 2;\n int a = nib_off ? (src[byte_off] >> 4) & 0x0F : src[byte_off] & 0x0F;\n uint8_t nib = (uint8_t)a & 0x0F;\n long d_byte_off = dst_off / 2;\n int d_nib_off = dst_off % 2;\n if (d_nib_off) {\n dst[d_byte_off] = (dst[d_byte_off] & 0x0F) | (nib << 4);\n } else {\n dst[d_byte_off] = (dst[d_byte_off] & 0xF0) | nib;\n }\n}\nCAST_IMPL(u4, u4)\n\n// Dispatch table\n\nstatic const cast_op_t cast_table[NX_NUM_DTYPES][NX_NUM_DTYPES] =\n {\n [NX_DTYPE_I8] =\n {\n [NX_DTYPE_I8] = nx_c_cast_i8_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_i8_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_i8_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_i8_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_i8_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_i8_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_i8_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_i8_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_i8_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_i8_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_i8_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_i8_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_i8_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_i8_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_i8_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_i8_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_i8_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_i8_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_i8_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_i8_to_f8e5m2,\n },\n [NX_DTYPE_U8] =\n {\n [NX_DTYPE_I8] = nx_c_cast_u8_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_u8_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_u8_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_u8_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_u8_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_u8_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_u8_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_u8_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_u8_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_u8_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_u8_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_u8_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_u8_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_u8_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_u8_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_u8_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_u8_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_u8_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_u8_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_u8_to_f8e5m2,\n },\n [NX_DTYPE_I16] =\n {\n [NX_DTYPE_I8] = nx_c_cast_i16_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_i16_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_i16_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_i16_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_i16_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_i16_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_i16_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_i16_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_i16_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_i16_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_i16_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_i16_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_i16_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_i16_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_i16_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_i16_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_i16_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_i16_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_i16_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_i16_to_f8e5m2,\n },\n [NX_DTYPE_U16] =\n {\n [NX_DTYPE_I8] = nx_c_cast_u16_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_u16_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_u16_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_u16_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_u16_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_u16_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_u16_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_u16_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_u16_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_u16_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_u16_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_u16_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_u16_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_u16_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_u16_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_u16_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_u16_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_u16_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_u16_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_u16_to_f8e5m2,\n },\n [NX_DTYPE_I32] =\n {\n [NX_DTYPE_I8] = nx_c_cast_i32_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_i32_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_i32_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_i32_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_i32_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_i32_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_i32_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_i32_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_i32_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_i32_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_i32_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_i32_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_i32_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_i32_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_i32_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_i32_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_i32_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_i32_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_i32_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_i32_to_f8e5m2,\n },\n [NX_DTYPE_I64] =\n {\n [NX_DTYPE_I8] = nx_c_cast_i64_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_i64_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_i64_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_i64_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_i64_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_i64_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_i64_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_i64_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_i64_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_i64_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_i64_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_i64_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_i64_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_i64_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_i64_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_i64_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_i64_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_i64_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_i64_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_i64_to_f8e5m2,\n },\n [NX_DTYPE_U32] =\n {\n [NX_DTYPE_I8] = nx_c_cast_u32_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_u32_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_u32_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_u32_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_u32_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_u32_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_u32_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_u32_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_u32_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_u32_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_u32_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_u32_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_u32_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_u32_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_u32_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_u32_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_u32_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_u32_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_u32_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_u32_to_f8e5m2,\n },\n [NX_DTYPE_U64] =\n {\n [NX_DTYPE_I8] = nx_c_cast_u64_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_u64_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_u64_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_u64_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_u64_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_u64_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_u64_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_u64_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_u64_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_u64_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_u64_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_u64_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_u64_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_u64_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_u64_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_u64_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_u64_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_u64_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_u64_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_u64_to_f8e5m2,\n },\n [NX_DTYPE_INAT] =\n {\n [NX_DTYPE_I8] = nx_c_cast_inat_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_inat_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_inat_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_inat_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_inat_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_inat_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_inat_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_inat_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_inat_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_inat_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_inat_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_inat_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_inat_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_inat_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_inat_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_inat_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_inat_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_inat_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_inat_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_inat_to_f8e5m2,\n },\n [NX_DTYPE_F16] =\n {\n [NX_DTYPE_I8] = nx_c_cast_f16_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_f16_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_f16_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_f16_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_f16_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_f16_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_f16_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_f16_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_f16_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_f16_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_f16_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_f16_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_f16_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_f16_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_f16_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_f16_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_f16_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_f16_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_f16_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_f16_to_f8e5m2,\n },\n [NX_DTYPE_F32] =\n {\n [NX_DTYPE_I8] = nx_c_cast_f32_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_f32_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_f32_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_f32_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_f32_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_f32_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_f32_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_f32_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_f32_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_f32_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_f32_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_f32_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_f32_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_f32_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_f32_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_f32_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_f32_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_f32_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_f32_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_f32_to_f8e5m2,\n },\n [NX_DTYPE_F64] =\n {\n [NX_DTYPE_I8] = nx_c_cast_f64_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_f64_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_f64_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_f64_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_f64_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_f64_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_f64_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_f64_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_f64_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_f64_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_f64_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_f64_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_f64_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_f64_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_f64_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_f64_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_f64_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_f64_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_f64_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_f64_to_f8e5m2,\n },\n [NX_DTYPE_C32] =\n {\n [NX_DTYPE_I8] = nx_c_cast_c32_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_c32_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_c32_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_c32_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_c32_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_c32_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_c32_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_c32_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_c32_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_c32_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_c32_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_c32_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_c32_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_c32_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_c32_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_c32_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_c32_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_c32_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_c32_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_c32_to_f8e5m2,\n },\n [NX_DTYPE_C64] =\n {\n [NX_DTYPE_I8] = nx_c_cast_c64_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_c64_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_c64_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_c64_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_c64_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_c64_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_c64_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_c64_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_c64_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_c64_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_c64_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_c64_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_c64_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_c64_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_c64_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_c64_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_c64_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_c64_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_c64_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_c64_to_f8e5m2,\n },\n [NX_DTYPE_BF16] =\n {\n [NX_DTYPE_I8] = nx_c_cast_bf16_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_bf16_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_bf16_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_bf16_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_bf16_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_bf16_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_bf16_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_bf16_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_bf16_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_bf16_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_bf16_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_bf16_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_bf16_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_bf16_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_bf16_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_bf16_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_bf16_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_bf16_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_bf16_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_bf16_to_f8e5m2,\n },\n [NX_DTYPE_BOOL] =\n {\n [NX_DTYPE_I8] = nx_c_cast_bool_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_bool_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_bool_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_bool_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_bool_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_bool_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_bool_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_bool_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_bool_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_bool_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_bool_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_bool_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_bool_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_bool_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_bool_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_bool_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_bool_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_bool_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_bool_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_bool_to_f8e5m2,\n },\n [NX_DTYPE_I4] =\n {\n [NX_DTYPE_I8] = nx_c_cast_i4_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_i4_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_i4_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_i4_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_i4_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_i4_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_i4_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_i4_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_i4_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_i4_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_i4_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_i4_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_i4_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_i4_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_i4_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_i4_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_i4_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_i4_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_i4_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_i4_to_f8e5m2,\n },\n [NX_DTYPE_U4] =\n {\n [NX_DTYPE_I8] = nx_c_cast_u4_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_u4_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_u4_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_u4_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_u4_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_u4_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_u4_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_u4_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_u4_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_u4_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_u4_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_u4_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_u4_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_u4_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_u4_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_u4_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_u4_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_u4_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_u4_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_u4_to_f8e5m2,\n },\n [NX_DTYPE_F8E4M3] =\n {\n [NX_DTYPE_I8] = nx_c_cast_f8e4m3_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_f8e4m3_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_f8e4m3_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_f8e4m3_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_f8e4m3_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_f8e4m3_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_f8e4m3_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_f8e4m3_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_f8e4m3_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_f8e4m3_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_f8e4m3_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_f8e4m3_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_f8e4m3_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_f8e4m3_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_f8e4m3_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_f8e4m3_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_f8e4m3_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_f8e4m3_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_f8e4m3_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_f8e4m3_to_f8e5m2,\n },\n [NX_DTYPE_F8E5M2] =\n {\n [NX_DTYPE_I8] = nx_c_cast_f8e5m2_to_i8,\n [NX_DTYPE_U8] = nx_c_cast_f8e5m2_to_u8,\n [NX_DTYPE_I16] = nx_c_cast_f8e5m2_to_i16,\n [NX_DTYPE_U16] = nx_c_cast_f8e5m2_to_u16,\n [NX_DTYPE_I32] = nx_c_cast_f8e5m2_to_i32,\n [NX_DTYPE_I64] = nx_c_cast_f8e5m2_to_i64,\n [NX_DTYPE_U32] = nx_c_cast_f8e5m2_to_u32,\n [NX_DTYPE_U64] = nx_c_cast_f8e5m2_to_u64,\n [NX_DTYPE_INAT] = nx_c_cast_f8e5m2_to_inat,\n [NX_DTYPE_F16] = nx_c_cast_f8e5m2_to_f16,\n [NX_DTYPE_F32] = nx_c_cast_f8e5m2_to_f32,\n [NX_DTYPE_F64] = nx_c_cast_f8e5m2_to_f64,\n [NX_DTYPE_C32] = nx_c_cast_f8e5m2_to_c32,\n [NX_DTYPE_C64] = nx_c_cast_f8e5m2_to_c64,\n [NX_DTYPE_BF16] = nx_c_cast_f8e5m2_to_bf16,\n [NX_DTYPE_BOOL] = nx_c_cast_f8e5m2_to_bool,\n [NX_DTYPE_I4] = nx_c_cast_f8e5m2_to_i4,\n [NX_DTYPE_U4] = nx_c_cast_f8e5m2_to_u4,\n [NX_DTYPE_F8E4M3] = nx_c_cast_f8e5m2_to_f8e4m3,\n [NX_DTYPE_F8E5M2] = nx_c_cast_f8e5m2_to_f8e5m2,\n }\n};\n\n// Dispatch function for cast operations\nstatic void dispatch_cast(value v_src, value v_dst) {\n ndarray_t src = extract_ndarray(v_src);\n ndarray_t dst = extract_ndarray(v_dst);\n\n if (src.ndim != dst.ndim) {\n cleanup_ndarray(&src);\n cleanup_ndarray(&dst);\n caml_failwith(\"shape mismatch\");\n }\n for (int i = 0; i < src.ndim; i++) {\n if (src.shape[i] != dst.shape[i]) {\n cleanup_ndarray(&src);\n cleanup_ndarray(&dst);\n caml_failwith(\"shape mismatch\");\n }\n }\n\n value v_src_data = Field(v_src, FFI_TENSOR_DATA);\n value v_dst_data = Field(v_dst, FFI_TENSOR_DATA);\n\n int src_kind = nx_buffer_get_kind(Caml_ba_array_val(v_src_data));\n int dst_kind = nx_buffer_get_kind(Caml_ba_array_val(v_dst_data));\n\n nx_dtype src_dtype = kind_to_dtype(src_kind);\n nx_dtype dst_dtype = kind_to_dtype(dst_kind);\n\n if (src_dtype == NX_NUM_DTYPES || dst_dtype == NX_NUM_DTYPES) {\n cleanup_ndarray(&src);\n cleanup_ndarray(&dst);\n caml_failwith(\"unsupported dtype\");\n }\n\n cast_op_t op = cast_table[src_dtype][dst_dtype];\n\n if (!op) {\n cleanup_ndarray(&src);\n cleanup_ndarray(&dst);\n caml_failwith(\"cast not supported for this dtype combination\");\n }\n\n caml_enter_blocking_section();\n op(&src, &dst);\n caml_leave_blocking_section();\n\n cleanup_ndarray(&src);\n cleanup_ndarray(&dst);\n}\n\n// OCaml FFI Stub\nCAMLprim value caml_nx_cast(value v_src, value v_dst) {\n CAMLparam2(v_src, v_dst);\n dispatch_cast(v_src, v_dst);\n CAMLreturn(Val_unit);\n}\n"
},
{
"path": "packages/nx/lib/backend_c/nx_c_cholesky.c",
"content": "/*---------------------------------------------------------------------------\n Copyright (c) 2026 The Raven authors. All rights reserved.\n SPDX-License-Identifier: ISC\n ---------------------------------------------------------------------------*/\n\n#include ![\"raven\"](\"www/site/raven.svg\")
\n![](\"figures/c_hubble.png\")
\n\n\n## Residuals: the acceleration signal\n\nResiduals $\\Delta\\mu = \\mu_\\text{obs} - \\mu_\\text{empty}(z)$ relative to an\nempty (coasting) universe isolate the acceleration signal. Positive residuals\nat high redshift mean supernovae are **fainter than expected** -- i.e. farther\naway than in a coasting universe. This is the direct evidence for cosmic\nacceleration.\n\nWe bin the data in redshift to show the trend clearly.\n\n\n\n```ocaml\nlet sn_mu_empty =\n let n = (Nx.shape sn_z).(0) in\n Nx.init f64 [| n |] (fun idx ->\n let z = Nx.scalar f64 (Nx.item [idx.(0)] sn_z) in\n Nx.item [] (Cosmo.distance_modulus ~p:p_empty z))\n\nlet sn_residual = Nx.sub sn_mu sn_mu_empty\n\n(* Model residuals on the grid *)\nlet res_lcdm = Nx.sub mu_lcdm mu_empty\nlet res_edsit = Nx.sub mu_edsit mu_empty\nlet res_open = Nx.sub mu_open mu_empty\n\n(* Bin the residuals using Talon grouping *)\nlet n_bins = 25\nlet log_z_min = Float.log10 0.01\nlet log_z_max = Float.log10 2.3\nlet bin_width = (log_z_max -. log_z_min) /. Float.of_int n_bins\n\nlet bin_df =\n let df = Talon.create [\n \"z\", Talon.Col.of_tensor sn_z;\n \"res\", Talon.Col.of_tensor sn_residual;\n ] in\n let df = Talon.filter_by df Talon.Row.(map (number \"z\") ~f:(fun z -> z > 0.01)) in\n Talon.with_column df \"bin\" f64 Talon.Row.(\n map (number \"z\") ~f:(fun z ->\n let b = int_of_float ((Float.log10 z -. log_z_min) /. bin_width) in\n Float.of_int (Int.max 0 (Int.min (n_bins - 1) b))))\n\nlet groups =\n Talon.group_by bin_df Talon.Row.(map (number \"bin\") ~f:int_of_float)\n |> List.filter (fun (_, g) -> Talon.num_rows g > 2)\n |> List.sort (fun (a, _) (b, _) -> Int.compare a b)\n\nlet n_groups = List.length groups\nlet bz = Nx.create f32 [| n_groups |]\n (Array.of_list (List.map (fun (_, g) -> Talon.Agg.mean g \"z\") groups))\nlet bmu = Nx.create f32 [| n_groups |]\n (Array.of_list (List.map (fun (_, g) -> Talon.Agg.mean g \"res\") groups))\nlet berr = Nx.create f32 [| n_groups |]\n (Array.of_list (List.map (fun (_, g) ->\n Talon.Agg.std g \"res\"\n /. Float.sqrt (Float.of_int (Talon.num_rows g - 1))) groups))\n```\n\n\nval sn_mu_empty : (float, Nx.float64_elt) Nx.t = float64 [1701] \n [28.5917, 28.5917, ..., 46.007, 46.5544]\nval sn_residual : (float, Nx.float64_elt) Nx.t = float64 [1701] \n [0.40697, 0.46417, ..., -0.583671, -0.371643]\nval res_lcdm : (float, Nx.float64_elt) Nx.t = float64 [200] \n [0.00189584, 0.00196019, ..., -0.320084, -0.343972]\nval res_edsit : (float, Nx.float64_elt) Nx.t = float64 [200] \n [-0.00170018, -0.00175822, ..., -1.13865, -1.16906]\nval res_open : (float, Nx.float64_elt) Nx.t = float64 [200] \n [-0.000498446, -0.000515476, ..., -0.465507, -0.480768]\nval n_bins : int = 25\nval log_z_min : float = -2.\nval log_z_max : float = 0.361727836017592841\nval bin_width : float = 0.094469113440703717\nval bin_df : Talon.t =\n \n\n![](\"figures/c_residual_plot.png\")
\n\n\n## Confidence contours in the $\\Omega_M$--$\\Omega_\\Lambda$ plane\n\nFollowing Perlmutter et al. (1999, Fig. 7), we scan a grid of\n$(\\Omega_M, \\Omega_\\Lambda)$ values and compute $\\chi^2$ at each point. The\nconfidence contours are drawn at $\\Delta\\chi^2 = 2.30, 6.17, 11.8$ (68.3%,\n95.4%, 99.7% for 2 parameters). We use only the Hubble-flow SNe ($z > 0.01$)\nand the diagonal errors (sufficient for this visualization).\n\n\n\n```ocaml\n(* Filter Hubble-flow SNe using Talon *)\nlet hf =\n Talon.filter_by df Talon.Row.(\n map2 (number \"zHD\") (number \"MU_SH0ES_ERR_DIAG\")\n ~f:(fun z err -> z > 0.01 && err > 0.0 && err < 10.0))\n\nlet hf_col name =\n Talon.get_column_exn hf name |> Talon.Col.to_tensor f64 |> Option.get\n\nlet hf_z = hf_col \"zHD\"\nlet hf_mu = hf_col \"MU_SH0ES\"\nlet hf_w = Nx.recip (Nx.square (hf_col \"MU_SH0ES_ERR_DIAG\"))\nlet n_hf = (Nx.shape hf_z).(0)\n\nlet () = Printf.printf \"Using %d Hubble-flow SNe for chi-squared grid\\n\" n_hf\n\n(* Chi-squared for a given (omega_m, omega_l) with M marginalized analytically.\n chi2 = sum w_i (mu_i - mu_th(z_i) - M)^2\n Minimizing over M: M* = sum(w_i * (mu_i - mu_th_i)) / sum(w_i)\n chi2_min = sum(w_i * d_i^2) - (sum(w_i * d_i))^2 / sum(w_i) *)\nlet hf_z_arr = Array.init n_hf (fun i -> Nx.item [i] hf_z)\nlet hf_mu_arr = Array.init n_hf (fun i -> Nx.item [i] hf_mu)\nlet hf_w_arr = Array.init n_hf (fun i -> Nx.item [i] hf_w)\nlet sum_w = Array.fold_left ( +. ) 0.0 hf_w_arr\n\n(* Pure-float distance modulus via 16-point Gauss-Legendre quadrature.\n Avoids all tensor allocation in the chi2 hot loop. *)\nlet gl_n = [| -0.9894009349916499; -0.9445750230732326; -0.8656312023878318;\n -0.7554044083550030; -0.6178762444026438; -0.4580167776572274;\n -0.2816035507792589; -0.0950125098376374; 0.0950125098376374;\n 0.2816035507792589; 0.4580167776572274; 0.6178762444026438;\n 0.7554044083550030; 0.8656312023878318; 0.9445750230732326;\n 0.9894009349916499 |]\nlet gl_wt = [| 0.0271524594117541; 0.0622535239386479; 0.0951585116824928;\n 0.1246289712555339; 0.1495959888165767; 0.1691565193950025;\n 0.1826034150449236; 0.1894506104550685; 0.1894506104550685;\n 0.1826034150449236; 0.1691565193950025; 0.1495959888165767;\n 0.1246289712555339; 0.0951585116824928; 0.0622535239386479;\n 0.0271524594117541 |]\n\nlet dist_mod_f omega_m omega_l z =\n let c_over_h0 = 299792.458 /. 70.0 in\n let omega_k = 1.0 -. omega_m -. omega_l in\n let half_z = z *. 0.5 in\n let integral = ref 0.0 in\n for k = 0 to 15 do\n let zp = half_z *. gl_n.(k) +. half_z in\n let opz = 1.0 +. zp in\n let ez = Float.sqrt (omega_m *. opz *. opz *. opz\n +. omega_k *. opz *. opz +. omega_l) in\n integral := !integral +. gl_wt.(k) /. ez\n done;\n let chi = c_over_h0 *. half_z *. !integral in\n let dl = (1.0 +. z) *. chi in\n 5.0 /. Float.log 10.0 *. Float.log dl +. 25.0\n\nlet chi2_at omega_m omega_l =\n let sum_wd = ref 0.0 in\n let sum_wdd = ref 0.0 in\n let ok = ref true in\n for i = 0 to n_hf - 1 do\n let mu_th_i = dist_mod_f omega_m omega_l hf_z_arr.(i) in\n if Float.is_nan mu_th_i then ok := false\n else begin\n let d = hf_mu_arr.(i) -. mu_th_i in\n let w = hf_w_arr.(i) in\n sum_wd := !sum_wd +. w *. d;\n sum_wdd := !sum_wdd +. w *. d *. d\n end\n done;\n if not !ok then infinity\n else !sum_wdd -. (!sum_wd *. !sum_wd /. sum_w)\n\n(* Scan the grid -- axis range matches Perlmutter 1999 Figure 7 *)\nlet n_om = 100\nlet n_ol = 100\nlet om_min = 0.0 and om_max = 3.0\nlet ol_min = -1.0 and ol_max = 3.0\n\nlet () = Printf.printf \"Computing chi-squared on %dx%d grid...\\n%!\" n_om n_ol\n\nlet chi2_grid =\n Nx.init f64 [| n_ol; n_om |] (fun idx ->\n let j = idx.(0) and i = idx.(1) in\n let omega_m = om_min +. (Float.of_int i +. 0.5) *. (om_max -. om_min) /. Float.of_int n_om in\n let omega_l = ol_min +. (Float.of_int j +. 0.5) *. (ol_max -. ol_min) /. Float.of_int n_ol in\n if omega_m < 0.001 then 1e10\n else chi2_at omega_m omega_l)\n\nlet chi2_min = Nx.item [] (Nx.min chi2_grid)\nlet delta_chi2 = Nx.sub_s chi2_grid chi2_min\n\nlet () =\n let flat_idx = Int32.to_int (Nx.item [] (Nx.argmin chi2_grid)) in\n let best_i = flat_idx mod n_om in\n let best_j = flat_idx / n_om in\n let best_om = om_min +. (Float.of_int best_i +. 0.5) *. (om_max -. om_min) /. Float.of_int n_om in\n let best_ol = ol_min +. (Float.of_int best_j +. 0.5) *. (ol_max -. ol_min) /. Float.of_int n_ol in\n Printf.printf \"Best fit: Omega_M = %.2f, Omega_Lambda = %.2f (chi2 = %.1f, dof ~ %d)\\n\"\n best_om best_ol chi2_min (n_hf - 1)\n```\n\n\nUsing 1590 Hubble-flow SNe for chi-squared grid\nComputing chi-squared on 100x100 grid...\nBest fit: Omega_M = 0.23, Omega_Lambda = 0.54 (chi2 = 684.2, dof ~ 1589)\nval hf : Talon.t =\n \n\n![](\"figures/c_contour.png\")
\n\n\n## Best-fit flat $\\Lambda$CDM\n\nRestricting to flat universes ($\\Omega_M + \\Omega_\\Lambda = 1$), we find the\nbest-fit $\\Omega_M$ by scanning along the flatness constraint and use Umbra's\n`Cosmo.flat_lcdm` to compute the corresponding distances.\n\n\n\n```ocaml\nlet n_flat = 200\nlet flat_om = Nx.linspace f64 0.01 0.99 n_flat\nlet flat_chi2 = Nx.init f64 [| n_flat |] (fun i ->\n let om = Nx.item [i.(0)] flat_om in\n chi2_at om (1.0 -. om))\n\nlet best_flat_i = Int32.to_int (Nx.item [] (Nx.argmin flat_chi2))\nlet omega_m_best = Nx.item [best_flat_i] flat_om\nlet omega_l_best = 1.0 -. omega_m_best\n\nlet () =\n Printf.printf \"\\n=== Flat ΛCDM best fit ===\\n\";\n Printf.printf \" Omega_M = %.3f\\n\" omega_m_best;\n Printf.printf \" Omega_L = %.3f\\n\" omega_l_best;\n Printf.printf \" chi2 = %.1f (dof ~ %d)\\n\"\n (Nx.item [best_flat_i] flat_chi2) (n_hf - 1);\n Printf.printf \"\\nPerlmutter et al. (1999) found Omega_M ~ 0.28, Omega_L ~ 0.72\\n\";\n Printf.printf \"Planck 2018 finds Omega_M = 0.315, Omega_L = 0.685\\n\"\n```\n\n\n\n=== Flat ΛCDM best fit ===\n Omega_M = 0.350\n Omega_L = 0.650\n chi2 = 684.6 (dof ~ 1589)\n\nPerlmutter et al. (1999) found Omega_M ~ 0.28, Omega_L ~ 0.72\nPlanck 2018 finds Omega_M = 0.315, Omega_L = 0.685\nval n_flat : int = 200\nval flat_om : (float, Nx.float64_elt) Nx.t = float64 [200] \n [0.01, 0.0149246, ..., 0.985075, 0.99]\nval flat_chi2 : (float, Nx.float64_elt) Nx.t = float64 [200] \n [1265.72, 1240.76, ..., 1313.88, 1321.44]\nval best_flat_i : int = 69\nval omega_m_best : float = 0.349798994974874378\nval omega_l_best : float = 0.650201005025125678\n\n\n### $\\chi^2$ profile along the flat-universe constraint\n\n\n\n```ocaml\nlet chi2_min_flat = Nx.item [best_flat_i] flat_chi2\nlet delta_flat = Nx.sub_s flat_chi2 chi2_min_flat\n\nlet _fig =\n Hugin.layers [\n Hugin.line ~x:(to32 flat_om) ~y:(to32 delta_flat)\n ~color:Hugin.Color.vermillion ~line_width:2.5 () ;\n Hugin.hline ~y:1.0 ~line_style:`Dashed ~color:Hugin.Color.gray\n ~label:\"Δχ² = 1 (1σ)\" () ;\n Hugin.hline ~y:4.0 ~line_style:`Dotted ~color:Hugin.Color.gray\n ~label:\"Δχ² = 4 (2σ)\" () ;\n Hugin.vline ~x:omega_m_best ~line_style:`Dashed\n ~color:Hugin.Color.sky_blue\n ~label:(Printf.sprintf \"Best fit: Ω_M = %.3f\" omega_m_best) () ;\n ]\n |> Hugin.xlim 0.0 1.0\n |> Hugin.ylim 0.0 20.0\n |> Hugin.xlabel \"Ω_M (flat universe)\"\n |> Hugin.ylabel \"Δχ²\"\n |> Hugin.title \"χ² Profile: Flat ΛCDM\"\n |> Hugin.legend ~loc:Hugin.Upper_right\n |> Hugin.grid_lines true\n```\n\n\nval chi2_min_flat : float = 684.600426847840708\nval delta_flat : (float, Nx.float64_elt) Nx.t = float64 [200] \n [581.117, 556.156, ..., 629.283, 636.844]\nval _fig : Hugin.t =\n \n\n![](\"figures/c_flat_chi2.png\")
\n\n\n## Cosmological implications\n\nWith the best-fit flat $\\Lambda$CDM parameters, we compute some fundamental\nproperties of the universe using Umbra's cosmology module.\n\n\n\n```ocaml\nlet p_best = Cosmo.flat_lcdm ~h0:70.0 ~omega_m:omega_m_best\n\nlet () =\n let z0 = Nx.scalar f64 0.0 in\n let z1 = Nx.scalar f64 1.0 in\n let z_star = Nx.scalar f64 1089.0 in\n Printf.printf \"\\n=== Universe properties (H₀ = 70 km/s/Mpc, Ω_M = %.3f) ===\\n\\n\" omega_m_best;\n Printf.printf \" Age of the universe = %.2f Gyr\\n\"\n (Nx.item [] (Unit.Time.in_gyr (Cosmo.age ~p:p_best z0)));\n Printf.printf \" Lookback time to z=1 = %.2f Gyr\\n\"\n (Nx.item [] (Unit.Time.in_gyr (Cosmo.lookback_time ~p:p_best z1)));\n Printf.printf \" Comoving distance to z=1 = %.0f Mpc\\n\"\n (Nx.item [] (Unit.Length.in_mpc (Cosmo.comoving_distance ~p:p_best z1)));\n Printf.printf \" Luminosity distance to z=1 = %.0f Mpc\\n\"\n (Nx.item [] (Unit.Length.in_mpc (Cosmo.luminosity_distance ~p:p_best z1)));\n Printf.printf \" Ang. diameter distance to z=1 = %.0f Mpc\\n\"\n (Nx.item [] (Unit.Length.in_mpc (Cosmo.angular_diameter_distance ~p:p_best z1)));\n Printf.printf \" Comoving distance to CMB = %.0f Mpc\\n\"\n (Nx.item [] (Unit.Length.in_mpc (Cosmo.comoving_distance ~p:p_best z_star)))\n```\n\n\n\n=== Universe properties (H₀ = 70 km/s/Mpc, Ω_M = 0.350) ===\n\n Age of the universe = 12.89 Gyr\n Lookback time to z=1 = 7.52 Gyr\n Comoving distance to z=1 = 3212 Mpc\n Luminosity distance to z=1 = 6423 Mpc\n Ang. diameter distance to z=1 = 1606 Mpc\n Comoving distance to CMB = 12758 Mpc\nval p_best : Umbra.Cosmo.params =