Add documentation for typescript via JSDoc (#45)

* Start JSDoc

* Rest of JS doc as things stand today

* CI config

* CI config

* Tweaks per comments

.github/workflows/docs.yaml

@@ -49,6 +49,16 @@ jobs:
           sudo apt-get update -y
           sudo apt-get install pandoc doxygen -y
 
+      - name: Set up Node
+        uses: actions/setup-node@v4
+        with:
+          node-version: 20
+          cache: 'yarn'
+          cache-dependency-path: clients/typescript/yarn.lock
+
+      - name: Install typescript dependencies
+        run: cd clients/typescript; yarn
+
       - name: Set up Julia
         uses: julia-actions/setup-julia@v2
 

clients/R/DESCRIPTION

@@ -8,7 +8,7 @@ Description: TinyStan provides the ability to use Stan's
     core sampling, optimization, and variational inference
     algorithms in a lightweight way.
 Encoding: UTF-8
-RoxygenNote: 7.3.1
+RoxygenNote: 7.3.2
 Roxygen: list(markdown = TRUE, r6 = TRUE)
 Suggests:
     testthat (>= 3.0.0)

clients/typescript/README.md

@@ -6,5 +6,4 @@ compiled with [emscripten](https://emscripten.org/).
 ## Example
 
 See the Stan Web Demo repository (https://github.com/WardBrian/stan-web-demo)
-for an example of how to use this package. More features and documentation
-will be added in the future.
+for an example of how to use this package.

clients/typescript/doc/jsdoc2md.json

@@ -0,0 +1,19 @@
+{
+  "source": {
+    "includePattern": ".+\\.ts(doc|x)?$",
+    "excludePattern": ".+\\.(test|spec).ts"
+  },
+  "template": "language-doc.hbs",
+  "plugins": ["plugins/markdown", "node_modules/jsdoc-babel", "@godaddy/dmd"],
+  "heading-depth": 3,
+  "babel": {
+    "extensions": ["ts", "tsx"],
+    "ignore": ["**/*.(test|spec).ts"],
+    "babelrc": false,
+    "presets": [
+      ["@babel/preset-env", { "targets": { "node": true } }],
+      "@babel/preset-typescript"
+    ],
+    "plugins": []
+  }
+}

clients/typescript/doc/language-doc.hbs

@@ -0,0 +1,44 @@
+<!-- This file is generated from the Handlebars template at `clients/typescript/doc/language-doc.hbs`
+
+If you are instead in `doc/language/js.md`, DO NOT EDIT THIS FILE DIRECTLY!
+-->
+
+# Typescript (JavaScript) Interface
+
+## Targetting WASM with TinyStan
+
+This interface is intended to be used to access models built
+with [Emscripten](https://emscripten.org/), which compiles C++ code to
+[WebAssembly](https://webassembly.org/).
+
+This requires some extra configuration compared to normal builds.
+Examples of the build (and usage) can be found in the following projects:
+
+- [Stan Web Demo](https://github.com/WardBrian/stan-web-demo/) ([Emscripten config](https://github.com/WardBrian/stan-web-demo/blob/main/build/local.mk))
+- [Stan Playground](https://github.com/flatironinstitute/stan-playground) ([Emscripten config](https://github.com/flatironinstitute/stan-playground/blob/main/backend/local.mk))
+
+Note that, among other things, a more recent version of TBB must be used,
+as WebAssembly is not supported in the version Stan currently vendors.
+
+## Installation
+
+The Typescript interface is [available on npm](https://www.npmjs.com/package/tinystan).
+To install it, run:
+```shell
+npm install tinystan --save
+```
+or
+```shell
+yarn add tinystan
+```
+
+<!-- hack to fix generated links
+     add more items here if sphinx warns about "undefined label"s
+-->
+[.sample(p)]:#stanmodel-sample-p
+[.pathfinder(p)]:#stanmodel-pathfinder-p
+[.stanVersion()]:#stanmodel-stanversion
+[.load(createModule, printCallback)]:#stanmodel-load-createmodule-printcallback
+
+## API Reference
+{{>main}}

clients/typescript/package.json

@@ -36,7 +36,8 @@
     "build": "tsup",
     "dev": "tsup --watch",
     "format": "prettier . --write",
-    "attw": "attw --pack ."
+    "attw": "attw --pack .",
+    "doc": "jsdoc2md --template ./doc/language-doc.hbs --plugin @godaddy/dmd --heading-depth=3 --configure ./doc/jsdoc2md.json --files src/*.ts "
   },
   "tsup": {
     "entry": [
@@ -59,14 +60,20 @@
   },
   "devDependencies": {
     "@arethetypeswrong/cli": "^0.15.3",
+    "@babel/cli": "^7.25.9",
+    "@babel/core": "^7.26.0",
+    "@babel/preset-env": "^7.26.0",
+    "@babel/preset-typescript": "^7.26.0",
+    "@godaddy/dmd": "^1.0.4",
     "@vitest/coverage-v8": "^1.6.0",
     "eslint": "^9.4.0",
+    "jsdoc-babel": "^0.5.0",
+    "jsdoc-to-markdown": "^9.0.5",
     "prettier": "^3.3.1",
     "tsup": "^8.1.0",
-    "typescript": "^5.2.2",
+    "typescript": "^5.6.3",
     "vitest": "^1.6.0"
   },
-  "dependencies": {},
   "publishConfig": {
     "access": "public"
   }

clients/typescript/src/model.ts

@@ -22,6 +22,11 @@ type model_ptr = internalTypes["model_ptr"];
 type error_ptr = internalTypes["error_ptr"];
 type cstr = internalTypes["cstr"];
 
+/**
+ * StanModel is a class that wraps the WASM module and provides a
+ * higher-level interface to the Stan library, abstracting away things
+ * like memory management and error handling.
+ */
 export default class StanModel {
   private m: WasmModule;
   private printCallback: PrintCallback | null;
@@ -34,6 +39,15 @@ export default class StanModel {
     this.sep = String.fromCharCode(m._tinystan_separator_char());
   }
 
+  /**
+   * Load a StanModel from a WASM module.
+   *
+   * @param {Function} createModule A function that resolves to a WASM module. This is
+   * much like the one Emscripten creates for you with `-sMODULARIZE`.
+   * @param {PrintCallback | null} printCallback A callback that will be called
+   * with any print statements from Stan. If null, this will default to `console.log`.
+   * @returns {Promise<StanModel>} A promise that resolves to a `StanModel`
+   */
   public static async load(
     createModule: (proto?: object) => Promise<WasmModule>,
     printCallback: PrintCallback | null,
@@ -81,7 +95,7 @@ export default class StanModel {
     }
   }
 
-  /**
+  /** @ignore
    * withModel serves as something akin to a context manager in
    * Python. It accepts the arguments needed to construct a model
    * (data and seed) and a callback.
@@ -116,6 +130,12 @@ export default class StanModel {
     }
   }
 
+  /**
+   * Sample using NUTS-HMC.
+   * @param {SamplerParams} p A (partially-specified) `SamplerParams` object.
+   * If a property is not specified, the default value will be used.
+   * @returns {StanDraws} A StanDraws object containing the parameter names and the draws
+   */
   public sample(p: Partial<SamplerParams>): StanDraws {
     const {
       data,
@@ -283,6 +303,13 @@ export default class StanModel {
     });
   }
 
+  /**
+   * Approximate the posterior using Pathfinder.
+   * @param {PathfinderParams} p A (partially-specified) `PathfinderParams` object.
+   * If a property is not specified, the default value will be used.
+   * @returns {StanDraws} A StanDraws object containing the parameter names and the
+   * approximate draws
+   */
   public pathfinder(p: Partial<PathfinderParams>): StanDraws {
     const {
       data,
@@ -388,6 +415,11 @@ export default class StanModel {
     });
   }
 
+  /**
+   * Get the version of the Stan library being used.
+   * @returns {string} The version of the Stan library being used,
+   * in the form "major.minor.patch"
+   */
   public stanVersion(): string {
     const major = this.m._malloc(4);
     const minor = this.m._malloc(4);

clients/typescript/src/types.ts

@@ -1,19 +1,90 @@
+/**
+ * @typedef {Record<string, unknown>} StanVariableInputs
+ * A type holding named inputs to a Stan model,
+ * e.g. the data or initial values.
+ */
 export type StanVariableInputs = Record<string, unknown>;
 
+/**
+ * @typedef {Function} PrintCallback
+ * A callback for printing output from the Stan model.
+ * @param {string} s The string to print.
+ * @returns {void}
+ */
 export type PrintCallback = (s: string) => void;
 
+/**
+ * The metric used for the HMC sampler.
+ * @enum {number}
+ * @readonly
+ * @property {0} UNIT - Unit metric.
+ * @property {1} DENSE - Dense metric.
+ * @property {2} DIAGONAL - Diagonal metric.
+ */
 export enum HMCMetric {
   UNIT = 0,
   DENSE = 1,
   DIAGONAL = 2,
 }
 
+/**
+ * @typedef {Object} StanDraws
+ * A type holding the result of a Stan sampling run.
+ * @property {string[]} paramNames The names of the parameters in
+ * the model
+ * @property {number[][]} draws A 2D array of draws from the posterior.
+ * The first dimension is the number of samples, and the second dimension
+ * is the number of parameters.
+ * @property {number[][] | number[][][] | undefined} metric The metric used for the
+ * HMC sampler. If the metric is not saved, this field is not present.
+ */
 export type StanDraws = {
   paramNames: string[];
   draws: number[][];
   metric?: number[][] | number[][][];
 };
 
+/**
+ * @typedef {Object} SamplerParams
+ * Parameters for the HMC sampler.
+ * @property {string | StanVariableInputs} [data=""] The data for the model
+ * @property {number} [num_chains=4] The number of chains to run
+ * @property {string | StanVariableInputs | string[] | StanVariableInputs[]} [inits=""]
+ * The initial values for the sampler. If an array, must have length `num_chains`.
+ * @property {number | null} [seed] The seed for the random number generator.
+ * If unspecified, a random seed will be generated.
+ * @property {number} [id=1] The ID for the first chain
+ * @property {number} [init_radius=2.0] Radius to initialize unspecified parameters within.
+ * The parameter values are drawn uniformly from the interval
+ * `[-init_radius, init_radius]` on the unconstrained scale.
+ * @property {number} [num_warmup=1000] The number of warmup iterations to run
+ * @property {number} [num_samples=1000] The number of samples to draw after warmup
+ * @property {HMCMetric} [metric=HMCMetric.DENSE] The type of mass matrix to use in the sampler
+ * @property {boolean} [save_metric=false] Whether to report the final mass matrix
+ * @property {number[] | number[][] | number[][][] | null} [init_inv_metric]
+ * The initial inverse metric to use. Currently, this argument is unused.
+ * @property {boolean} [adapt=true] Whether the sampler should adapt the step size and metric
+ * @property {number} [delta=0.8] Target acceptance rate
+ * @property {number} [gamma=0.05] Adaptation regularization scale
+ * @property {number} [kappa=0.75] Adaptation relaxation exponent
+ * @property {number} [t0=10.0] Adaptation iteration offset
+ * @property {number} [init_buffer=75] Number of warmup samples to use for initial
+ * step size adaptation.
+ * @property {number} [term_buffer=50] Number of warmup samples to use for step size
+ * adaptation after the metric is adapted
+ * @property {number} [window=25] Initial number of iterations to use for metric adaptation,
+ * which is doubled each time the adaptation window is hit
+ * @property {boolean} [save_warmup=false] Whether to save the warmup draws
+ * @property {number} [stepsize=1.0] Initial step size for the sampler
+ * @property {number} [stepsize_jitter=0.0] Amount of random jitter to add to the step size
+ * @property {number} [max_depth=10] Maximum tree depth for the NUTS sampler
+ * @property {number} [refresh=0] Number of iterations between progress messages.
+ * If 0, no output is printed.
+ * @property {number} [num_threads=-1] Number of threads to use for sampling.
+ * If -1, the number of threads is determined by the number of available CPU cores.
+ * May not be supported in all environments, and requires extra configuration
+ * during the Emscripten compilation.
+ */
 export interface SamplerParams {
   data: string | StanVariableInputs;
   num_chains: number;
@@ -69,6 +140,44 @@ interface PathfinderUniqueParams {
   num_threads: number;
 }
 
+/**
+ * @typedef {Object} PathfinderParams
+ * Parameters for the Pathfinder algorithm.
+ * @property {string | StanVariableInputs} [data=""] The data for the model
+ * @property {number} [num_paths=4] The number of individual paths to run
+ * @property {string | StanVariableInputs | string[] | StanVariableInputs[]} [inits=""]
+ * The initial values for the algorithm. If an array, must have length `num_paths`.
+ * @property {number | null} [seed] The seed for the random number generator.
+ * If unspecified, a random seed will be generated.
+ * @property {number} [id=1] The ID for the first path
+ * @property {number} [init_radius=2.0] Radius to initialize unspecified parameters within.
+ * The parameter values are drawn uniformly from the interval
+ * `[-init_radius, init_radius]` on the unconstrained scale.
+ * @property {number} [num_draws=1000] The number of draws to take for each path
+ * @property {number} [max_history_size=5] History size used by the internal L-BFGS algorithm
+ * to approximate the Hessian
+ * @property {number} [init_alpha=0.001] Initial step size for the internal L-BFGS algorithm
+ * @property {number} [tol_obj=1e-12] Convergence tolerance for the objective function
+ * @property {number} [tol_rel_obj=1e4] Relative convergence tolerance for the objective function
+ * @property {number} [tol_grad=1e-8] Convergence tolerance for the gradient norm
+ * @property {number} [tol_rel_grad=1e7] Relative convergence tolerance for the gradient norm
+ * @property {number} [tol_param=1e-8] Convergence tolerance for the changes in parameters
+ * @property {number} [num_iterations=1000] Maximum number of iterations for the internal
+ * L-BFGS algorithm
+ * @property {number} [num_elbo_draws=25] Number of Monte Carlo draws used to estimate the ELBO
+ * @property {number} [num_multi_draws=1000] Number of draws returned by Multi-Pathfinder
+ * @property {boolean} [calculate_lp=true] Whether to calculate the log probability of the
+ * approximate draws.
+ * If false, this also implies `psis_resample=false`.
+ * @property {boolean} [psis_resample=true] Whether to use Pareto smoothed importance sampling on
+ * the approximate draws. If false, all `num_paths * num_draws` approximate samples will be returned.
+ * @property {number} [refresh=0] Number of iterations between progress messages.
+ * If 0, no output is printed.
+ * @property {number} [num_threads=-1] Number of threads to use for Pathfinder.
+ * If -1, the number of threads is determined by the number of available CPU cores.
+ * May not be supported in all environments, and requires extra configuration
+ * during the Emscripten compilation.
+ */
 export type PathfinderParams = LBFGSConfig & PathfinderUniqueParams;
 
 // ------------- internal types -------------