madspace: audit comments and docstrings against the source

madspace/generate_pyi.py

@@ -2,8 +2,8 @@
 """Generate madspace/_madspace_py.pyi from a freshly built _madspace_py module.
 
 pybind11-stubgen has to import the module to introspect it, which in turn
-requires the full "madspace" package (its __init__.py preloads libmadspace
-and monkey-patches a few classes). Rather than requiring a full install, this
+requires the full "madspace" package (its loader preloads libmadspace and
+monkey-patches a few classes). Rather than requiring a full install, this
 script assembles a throwaway copy of just the pieces needed for that import
 in --stage-dir, then runs pybind11-stubgen against it.
 """
@@ -77,9 +77,8 @@ def main():
     )
 
     # pybind11-stubgen's own error checks (above) don't guarantee the emitted
-    # file is syntactically valid Python, e.g. a bad --print-safe-value-reprs
-    # match can inline an expression verbatim; parse it so a broken stub
-    # fails the build instead of shipping something downstream tools choke on.
+    # file is syntactically valid Python (e.g. an expression may be inlined
+    # verbatim); parse it so a broken stub fails the build.
     stub_path = args.output_dir / "madspace" / "_madspace_py.pyi"
     source = stub_path.read_text()
     try:

madspace/include/madspace/compgraphs/function.hpp

@@ -31,7 +31,7 @@ struct InstructionCall {
 /**
  * An immutable compute graph produced by @ref FunctionBuilder.
  *
- * A function has named `inputs` and `outputs`, a list of intermediate `locals`,
+ * A function has named `inputs` and `outputs`, a list of `locals`,
  * named `globals` for trainable or externally supplied tensors, and the ordered
  * list of @ref InstructionCall that computes the outputs from the inputs. Run
  * one with a `FunctionRuntime`, or `save` it to disk and `load` it back.
@@ -46,7 +46,7 @@ class Function {
     const NamedVector<Value>& inputs() const { return _inputs; }
     /// Named graph outputs.
     const NamedVector<Value>& outputs() const { return _outputs; }
-    /// Intermediate values, one per instruction output.
+    /// All local values: inputs, globals, literals and instruction outputs.
     const ValueVec& locals() const { return _locals; }
     /// Named global tensors: trainable parameters and external inputs.
     const std::vector<std::pair<std::string, Value>>& globals() const {
@@ -97,7 +97,7 @@ void from_json(const nlohmann::json& j, Function& call);
  *
  * Construct it with the input and output types, read the graph inputs with
  * `input`, add nodes by calling the instruction methods (for example `add`,
- * `stack` or `two_body_decay_com`, listed below), wire the results to the
+ * `stack` or `two_body_decay_com`, generated below), wire the results to the
  * outputs with `output`, then call `function` to obtain the finished
  * @ref Function. The instruction methods are generated from
  * `madspace/instruction_set.yaml`.
@@ -108,7 +108,8 @@ class FunctionBuilder {
     FunctionBuilder(
         const NamedVector<Type>& _input_types, const NamedVector<Type>& _output_types
     );
-    /// Start a builder pre-populated with an existing function.
+    /// Start a builder that has the inputs of an existing function; its
+    /// instructions are not copied.
     FunctionBuilder(const Function& function);
     /// The graph input at position @p index.
     Value input(int index) const;
@@ -134,7 +135,8 @@ class FunctionBuilder {
     void set_current_stream(std::size_t stream_index) {
         _current_stream = stream_index;
     }
-    /// Freeze the recorded graph into an immutable @ref Function.
+    /// Freeze the recorded graph into an immutable @ref Function. Instructions
+    /// that do not contribute to any output are dropped.
     Function function();
 
     /// Sum of @p values, element by element.

madspace/include/madspace/compgraphs/type.hpp

@@ -72,7 +72,7 @@ class BatchSize {
     BatchSize operator-(const BatchSize& other) const { return add(other, -1); }
     bool operator==(const BatchSize& other) const { return value == other.value; }
     bool operator!=(const BatchSize& other) const { return value != other.value; }
-    /// Human-readable representation such as `"batch_size + 1"`.
+    /// Human-readable representation such as `"batch_size+1"`.
     std::string to_string() const;
 
     friend std::ostream& operator<<(std::ostream& out, const BatchSize& batch_size);
@@ -187,16 +187,17 @@ using LiteralValue = std::variant<me_int_t, double, TensorValue, std::monostate>
  * or a reference to a graph local.
  *
  * @ref FunctionBuilder methods take and return `Value`s. A value with
- * `local_index >= 0` refers to the output of an earlier instruction. Otherwise
+ * `local_index >= 0` refers to a local of the function being built (an input,
+ * global, registered literal or instruction output). Otherwise
  * @p literal_value holds an inline scalar or dense tensor constant. Scalars and
- * nested `std::vector`s convert to a `Value` implicitly.
+ * `std::vector`s (also nested ones) convert to a `Value` implicitly.
  */
 struct Value {
     /// Type of the value.
     Type type;
-    /// Inline constant, or `std::monostate` for a graph local.
+    /// Inline constant, or `std::monostate` for a graph local without constant value.
     LiteralValue literal_value;
-    /// Index of the referenced graph local, or -1 for a literal.
+    /// Index of the referenced graph local, or -1 if not yet registered as one.
     int local_index = -1;
 
     /// Empty value (no literal, no local); tests as `false`.

madspace/include/madspace/driver/adam_optimizer.hpp

@@ -6,9 +6,9 @@
 
 namespace madspace {
 
-/// Compute-graph function that rescales a gradient tensor so its norm never
-/// exceeds a threshold; the graph built internally by @ref AdamOptimizer for
-/// `grad_clip_threshold`.
+/// Compute-graph function that rescales `gradients_in` so that its L2 norm does
+/// not exceed `threshold`; also returns the unclipped norm. Used internally by
+/// @ref AdamOptimizer for `grad_clip_threshold`.
 class GradientClipper : public FunctionGenerator {
 public:
     GradientClipper();
@@ -27,7 +27,9 @@ class GradientClipper : public FunctionGenerator {
  * decay, @p eps for numerical stability, optional @p weight_decay) to its
  * `requires_grad` @ref Context globals. @p schedule optionally decays the
  * learning rate over @p step_count steps; @p grad_clip_threshold, if
- * non-zero, rescales the gradient norm before the update.
+ * non-zero, clips the gradient norm to `grad_clip_threshold` times the square
+ * root of the number of trained parameters before the update. A step is skipped
+ * if the loss is NaN or, after 100 steps, exceeds 20 times its running mean.
  *
  * **References**
  * - [1] D. P. Kingma, J. Ba, "Adam: A method for stochastic optimization",
@@ -54,8 +56,10 @@ class AdamOptimizer {
      * @param beta1               First-moment decay rate.
      * @param beta2               Second-moment decay rate.
      * @param eps                 Numerical-stability constant.
-     * @param grad_clip_threshold Maximum gradient norm; `0` disables clipping.
-     * @param weight_decay        L2 weight-decay coefficient.
+     * @param grad_clip_threshold Maximum gradient norm per sqrt(parameter
+     *                            count); `0` disables clipping.
+     * @param weight_decay        Decoupled weight decay (parameters shrink
+     *                            by `learning_rate * weight_decay` per step).
      */
     AdamOptimizer(
         const Function& function,
@@ -69,8 +73,8 @@ class AdamOptimizer {
         double grad_clip_threshold = 0.0,
         double weight_decay = 0.0
     );
-    /// Run one training step on `inputs` and apply the Adam update; returns
-    /// `function`'s outputs.
+    /// Run one training step on `inputs` and apply the Adam update (unless the
+    /// step is skipped, see above); returns `function`'s outputs.
     TensorVec step(const TensorVec& inputs);
     /// Replace the loss function, keeping the optimizer state.
     void replace_function(const Function& function);

madspace/include/madspace/driver/backend.hpp

@@ -36,11 +36,12 @@ class Runtime {
 /// Owning handle to a `Runtime`.
 using RuntimePtr = std::unique_ptr<Runtime>;
 /// Compile `function` into a callable @ref Runtime for `context`'s device.
-/// `concurrent` allows several `Runtime`s to share one compilation cache.
+/// `concurrent` lets the CPU runtime run independent instructions concurrently on
+/// the context's thread pool; the GPU runtime ignores it.
 RuntimePtr
 build_runtime(const Function& function, ContextPtr context, bool concurrent = true);
-/// Names of the compute backends (CPU, CUDA, HIP, ...) built into this
-/// installation.
+/// Names of the CPU backend variants usable on this host ("cpu_scalar",
+/// "cpu_vec128", "cpu_vec256", ...); GPU backends are not listed.
 std::vector<std::string> available_backends();
 DevicePtr cpu_device();
 DevicePtr cuda_device(std::size_t index = 0);

madspace/include/madspace/driver/channel_generator.hpp

@@ -31,8 +31,8 @@ namespace madspace {
  */
 class ChannelEventGenerator {
 public:
-    /// Reconstruct a generator previously written with @ref save, reopening
-    /// `event_file` and `weight_file`.
+    /// Reconstruct a generator previously written with @ref save. Creates new,
+    /// empty `event_file` and `weight_file`.
     static ChannelEventGenerator load(
         const std::string& channel_file,
         const std::vector<ContextPtr>& contexts,
@@ -53,8 +53,9 @@ class ChannelEventGenerator {
     /**
      * @param contexts          One context per device to run on.
      * @param integrand         The channel's compiled integrand.
-     * @param event_file        Path of the weighted event file.
-     * @param weight_file       Path of the max-weight tracking file.
+     * @param event_file        Path of the file the unweighted events are
+     *                          written to.
+     * @param weight_file       Path of the file holding one weight per event.
      * @param config            Generator configuration.
      * @param subprocess_index  Index of the owning subprocess.
      * @param name              Human-readable channel name.
@@ -87,9 +88,9 @@ class ChannelEventGenerator {
     std::size_t count_requested_opt() const { return _count_requested_opt; }
     /// The filled observable histograms, if any were requested.
     const std::vector<Histogram>& histograms() const { return _histograms; }
-    /// The weighted event file.
+    /// The event file.
     EventFile& event_file() { return _event_file; }
-    /// The max-weight tracking file.
+    /// The weight file, one weight per event of @ref event_file.
     EventFile& weight_file() { return _weight_file; }
     /// Current maximum-weight estimate, used for unweighting.
     double max_weight() const { return _max_weight; }
@@ -103,8 +104,8 @@ class ChannelEventGenerator {
     bool has_abs_integral_prior() const { return _abs_integral_prior.has_value(); }
     /// Set the signed integral known from an earlier run, with its standard error.
     void set_integral_prior(double mean, double error);
-    /// Whether this run has too few samples for an integral estimate of its own, so
-    /// that @ref EventGenerator uses the priors instead.
+    /// Whether an absolute integral prior is set and this run has too few samples
+    /// for an estimate of its own, so that @ref EventGenerator uses the priors.
     bool uses_integral_prior() const;
     /// The prior of the signed integral as (mean, standard error), or of the
     /// absolute one if no signed prior was set.
@@ -117,7 +118,7 @@ class ChannelEventGenerator {
     /// with the prior (if any) weighted by the inverse variances. Used to share
     /// the events between channels.
     double abs_integral_estimate() const;
-    /// Current generation batch size.
+    /// Size of the next optimization batch, see @ref next_vegas_batch_size.
     std::size_t batch_size() const { return _batch_size; }
     /// Whether a VEGAS grid or channel-weight optimization pass is still
     /// pending.
@@ -138,16 +139,17 @@ class ChannelEventGenerator {
     /// Layout of @ref event_file.
     const DataLayout& event_file_layout() const { return _event_file_layout; }
 
-    /// Run a final unweighting pass over @ref weight_file into @ref
-    /// event_file, using `rand_gen`.
+    /// Unweight @ref weight_file in place with the current maximum weight: weights
+    /// of rejected events are set to zero. Uses `rand_gen`.
     void unweight_file(MixMaxRandom& rand_gen);
-    /// Integrate `job`'s events, accumulating @ref cross_section.
+    /// Accumulate `job`'s weights into @ref cross_section, the status and the
+    /// optimizer and histogram data.
     void integrate(const GeneratorBatchJob& job);
-    /// Accumulate `job`'s events into the VEGAS grid and channel-weight
-    /// optimizers.
+    /// Update the VEGAS grid and channel-weight optimizers from the data
+    /// accumulated by @ref integrate, and update the optimization status.
     void optimize_vegas(const GeneratorBatchJob& job);
-    /// Sum of this channel's weight over `event_count` freshly sampled events,
-    /// used for multi-channel weight optimization.
+    /// Sum of |weight| / max_weight over the first `event_count` nonzero weights
+    /// of @ref weight_file.
     double channel_weight_sum(std::size_t event_count);
     /// Dispatch `job` (a survey or generation batch) onto the thread pool,
     /// seeded from `seed`.
@@ -168,20 +170,22 @@ class ChannelEventGenerator {
     void start_unweight_job(GeneratorBatchJob& job, ResultQueue& result_queue);
     /// Size of the next VEGAS/channel-weight optimization batch.
     std::size_t next_vegas_batch_size();
-    /// Discard the accumulated event and weight files.
+    /// Discard the accumulated events and weights, and reset the integral, the
+    /// event counts and (unless fixed) the maximum weight.
     void clear_events();
     /// Update @ref max_weight from a batch of `weights`.
     void update_max_weight(Tensor weights);
-    /// Append `unweighted_events` to @ref event_file, recording
-    /// `job_max_weight`.
+    /// Append `unweighted_events` to @ref event_file and their weights to
+    /// @ref weight_file. `job_max_weight` is the maximum weight they were
+    /// unweighted with.
     void write_events(const TensorVec& unweighted_events, double job_max_weight);
     /// Serialize this generator's state to `file_name`; see @ref load.
     void save(const std::string& file_name) const;
     /// The JSON text that @ref save writes; see @ref load_json. With
-    /// `include_estimates`, the current maximum weight and the integral of the
-    /// absolute weights are stored as well. @ref load_json then keeps the
-    /// maximum weight fixed and uses the integral as a prior, see
-    /// @ref abs_integral_estimate.
+    /// `include_estimates`, the current maximum weight and, if there are at least
+    /// two samples, the integrals (signed and absolute) are stored as well.
+    /// @ref load_json then keeps the maximum weight fixed and uses the
+    /// integrals as priors, see @ref abs_integral_estimate.
     std::string to_json(bool include_estimates = false) const;
 
 private:
@@ -251,10 +255,9 @@ class ChannelEventGenerator {
     RunningIntegral _cross_section;
     RunningIntegral _abs_cross_section;
     double _max_weight = 0.;
-    // Expected number of accepted events, sum of min(|w| / max_weight, 1) over the
-    // committed generate stages, rescaled like count_unweighted when the max
-    // weight rises. Decides when the max weight is frozen: unlike count_unweighted
-    // it only depends on the generate stages, which commit in a fixed order.
+    // Sum of min(|w| / max_weight, 1) over the committed generate stages, rescaled
+    // when the max weight rises. Decides when the max weight is frozen; unlike
+    // count_unweighted it does not depend on the unweight stages' commit timing.
     double _expected_unweighted = 0.;
     bool _max_weight_fixed = false;
     // mean and standard error
@@ -265,11 +268,8 @@ class ChannelEventGenerator {
     // never reset. Separate for survey/generate so neither depends on the other.
     std::size_t _survey_rng_seq = 0;
     std::size_t _generate_rng_seq = 0;
-    // Progress of unweight_file()'s final pass over _weight_file: _unweighted_count
-    // is the index up to which it's been scanned under the *current* _max_weight,
-    // and _unweighted_accept_count the number of those accepted. Both reset to 0
-    // whenever _max_weight changes (forcing a full rescan), since a changed
-    // max_weight invalidates every prior accept/reject decision.
+    // Progress of unweight_file(): the number of weights scanned and accepted under
+    // the current _max_weight. Reset to 0 whenever _max_weight changes.
     std::size_t _unweighted_count = 0;
     std::size_t _unweighted_accept_count = 0;
     std::size_t _iters_without_improvement = 0;