oxieml 0.1.2

EML operator: all elementary functions from exp(x) - ln(y)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

//! Post-Adam constant rounding and named-constant extraction.
//!
//! After Adam optimisation converges, we optionally snap free constants to:
//! 1. **Integer rounding** — parameters within 0.02 of an integer are rounded.
//!    Controlled by [`SymRegConfig::integer_rounding`].
//! 2. **Named-constant extraction** — candidate set: π, e, √2, simple rationals
//!    with denominator ≤ 12 (Stern-Brocot). Acceptance criterion:
//!    `new_mse ≤ (1 + eps) * current_mse`.
//!    Controlled by [`SymRegConfig::constant_extraction`].

use crate::grad::ParameterizedEmlTree;
use crate::lower::LoweredOp;
use crate::tree::EmlTree;

use super::constants::{bake_params_into_lowered, extract_named_constants};
use super::topology::{compute_mse_parameterized, try_integer_rounding};

/// Apply integer rounding to `best_params` and accept if MSE stays within 1%.
///
/// Returns the updated `(params, mse)` pair. If rounding degrades MSE by
/// more than 1%, the original values are returned unchanged.
pub(super) fn try_post_adam_rounding(
    topology: &EmlTree,
    best_params: Vec<f64>,
    best_mse: f64,
    inputs: &[Vec<f64>],
    targets: &[f64],
) -> (Vec<f64>, f64) {
    let rounded = try_integer_rounding(&best_params);
    let mut ptree_rounded = ParameterizedEmlTree::from_topology(topology, 1.0);
    ptree_rounded.params = rounded;
    let rounded_mse = compute_mse_parameterized(&ptree_rounded, inputs, targets);
    if let Some(rmse) = rounded_mse {
        if rmse <= best_mse * 1.01 {
            return (ptree_rounded.params, rmse);
        }
    }
    (best_params, best_mse)
}

/// Bake learned parameters into the lowered form and optionally extract named
/// constants (π, e, √2, simple rationals).
///
/// Returns `(final_lowered_op, final_mse)`. If `constant_extraction` is `None`,
/// the baked-but-not-extracted form is returned.
pub(super) fn try_extract_named_constants(
    topology: &EmlTree,
    best_params: &[f64],
    best_mse: f64,
    constant_extraction: Option<f64>,
    inputs: &[Vec<f64>],
    targets: &[f64],
) -> (LoweredOp, f64) {
    let baked = bake_params_into_lowered(topology, best_params);
    let baked_simplified = baked.simplify();

    if let Some(eps) = constant_extraction {
        extract_named_constants(baked_simplified, best_mse, eps, inputs, targets)
    } else {
        (baked_simplified, best_mse)
    }
}