Skip to content

Refactor approximations into a pure-JuMP layer #101

Closed
acostarelli wants to merge 10 commits into
mainfrom
ac/approx-pure-jump-refactor
Closed

Refactor approximations into a pure-JuMP layer #101
acostarelli wants to merge 10 commits into
mainfrom
ac/approx-pure-jump-refactor

Conversation

@acostarelli
Copy link
Copy Markdown
Member

For the purpose of more easily testing approximations separate from container API. Also includes speedups by using vectorized JuMP object creation. Potentially useful for Data Centers work.

…ner API

Previously the quadratic and bilinear approximation builders interleaved
JuMP-construction with OptimizationContainer bookkeeping, making the math
hard to test independently. This commit splits each approximation into two
layers:

- A pure-JuMP `build_quadratic_approx` / `build_bilinear_approx` per method
  that operates on a bare `JuMP.Model`, JuMP containers, and bounds, and
  returns a result struct holding all the JuMP objects (variables,
  constraints, expressions).

- A generic IOM wrapper in `approximations/common.jl` that preserves the
  existing `_add_quadratic_approx!` / `_add_bilinear_approx!` POM entry
  points: it calls `build_*` and then dispatches `register_in_container!`
  on the result struct to write all auxiliary JuMP objects into the
  container with the right keys and meta suffix.

The directories `quadratic_approximations/` and `bilinear_approximations/`
are merged into a single `approximations/` folder with one self-contained
file per method (config, result, build, register all colocated). The math
layer is now exercisable without any container scaffolding — `test/
test_pure_jump_approximations.jl` does exactly this.

Optional add-ons (PWMCC concave cuts, epigraph tightening, reformulated
McCormick) now live inside the corresponding `build_*` function rather
than as separate post-processing calls; result structs carry them as
optional fields and the register dispatch propagates them appropriately.

Loops over `(name, t)` are replaced with batched `JuMP.@variable`,
`JuMP.@constraint`, and `JuMP.@expression` macros where applicable for
both readability and a modest speedup.

POM-facing signatures are unchanged. All 1098 IOM tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@acostarelli acostarelli requested a review from annacasavant May 13, 2026 22:06
@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented May 13, 2026

Performance Results
Main


This branch


Copy link
Copy Markdown
Member

@jd-lara jd-lara left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the rationale for this change?

Copy link
Copy Markdown
Member Author

@acostarelli acostarelli left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review 1

Overall, don't take this as just nit-picks. I think the major refactor is not done yet. The jump layer still looks like it did with containers just now we're calling them dense axis arrays. I want to be able to read the jump code and cleanly see the formulation without other noise. The name and time axes are fluff that we handle after.

Claude:

Resolved

  • Vectorized every register_in_container! copy looptarget.data .= source.data on the underlying Array, no more per-(name, k, t) loops.
  • Typed every *Result struct's parameters with concrete bounds (<: QuadraticApproxResult, <: DenseAxisArray{AffExpr, 2}, etc.).
  • Removed sparse = true from src/approximations/. Split the generic McCormick into McCormickLowerConstraint / McCormickUpperConstraint so the lower side cleanly disappears when lower_bounds = false; densely-filled containers (HybSBoundConstraint, reformulated McCormick, sawtooth MIP, epigraph tangent) all switched to dense. Also fixed a latent sawtooth schema bug — adjacency container was declared 3D but populated 4D.
  • Renamed _add_quadratic_approx!add_quadratic_approx! and _add_bilinear_approx!add_bilinear_approx!. POM call sites need the paired rename.
  • Converted !== nothing guards to Nothing dispatch for register_reformulated_mccormick!, register_mccormick_envelope!, _register_tightening!, register_pwmcc!, plus a new _register_sawtooth_tightening!.
  • Dropped redundant MinMax(...) wrapper (bin2.jl, hybs.jl) and collect(name_axis) throughout.
  • Vectorized JuMP-side construction loops: pwmcc_cuts.jl brk, sawtooth MIP (4 vectorized @constraint families instead of a triple loop), epigraph LP and tangent cuts, manual-SOS2 adjacency, and build_mccormick_envelope.
  • Extracted scale_back_g_basis(...) in epigraph.jl and reused it from sawtooth and the epigraph tangent cuts.
  • Restored precomputed-form add_bilinear_approx! entrypoints for every bilinear method (Bin2, HybS, NMDT, DNMDT, NoBilinear). Backed by a small _PrebuiltQuadApprox adapter in common.jl.
  • pwmcc_cuts.jl:144: confirmed JuMP's macro parser recognizes the sum(...) pattern.

Deferred (next round)

  • The architectural ask itself — split JuMP math into a single-(name, t)-element routine and broadcast axes on top: bin2.jl:41, mccormick.jl:71, nmdt_discretization.jl:265, plus the abstract-Result base at nmdt_bilinear.jl:242, register-side constraint organization at epigraph.jl:138, common-check wrapper at hybs.jl:59/:123, and loop-count reduction at nmdt_discretization.jl:388.
  • epigraph.jl:249 (numeric axis → string axis): downstream to_output_dataframe behavior, skipped for now.

All 1098 IOM tests pass.

Comment thread src/approximations/bin2.jl Outdated
Comment thread src/approximations/bin2.jl Outdated
Comment thread src/approximations/bin2.jl Outdated
Comment thread src/approximations/common.jl Outdated
Comment thread src/approximations/common.jl Outdated
Comment thread src/approximations/nmdt_discretization.jl Outdated
Comment thread src/approximations/nmdt_quadratic.jl Outdated
Comment thread src/approximations/pwmcc_cuts.jl Outdated
Comment thread src/approximations/pwmcc_cuts.jl Outdated
Comment thread src/approximations/sawtooth.jl Outdated
Anthony Costarelli and others added 9 commits May 14, 2026 17:16
Targeted at the smaller, non-architectural review threads on PR #101.

Cross-cutting changes:
- Vectorize every `register_in_container!` copy loop: replace per-element
  loops with `target.data .= source.data` broadcasts on the underlying
  Arrays, since `add_*_container!` returns `DenseAxisArray`s whose internal
  `.data` matches the build-side container's shape.
- Type every approximation `*Result` struct's parameters with concrete
  bounds (e.g. `XSQ <: QuadraticApproxResult`, `A <: DenseAxisArray{AffExpr, 2}`)
  rather than bare typevars.
- Drop `sparse = true` everywhere in `src/approximations/`. Split the
  generic McCormick envelope key into `McCormickLowerConstraint` /
  `McCormickUpperConstraint` so the lower side cleanly disappears when
  `lower_bounds = false` (NMDT's `tighten = true` path). Densely-populated
  containers (`HybSBoundConstraint`, reformulated McCormick, sawtooth MIP)
  switch to dense.
- Drop `collect(name_axis)` calls; `add_*_container!` accepts any iterable
  as an axis, and the source axis is already a `Vector{String}` in
  practice.

Item-specific changes:
- Rename `_add_quadratic_approx!` → `add_quadratic_approx!` and
  `_add_bilinear_approx!` → `add_bilinear_approx!`; POM call sites need a
  paired update.
- Convert `if x !== nothing; register_*!(...); end` guards to dispatched
  no-op `register_*!(_, _, ::Nothing, _)` overloads. Same for
  `_register_tightening!`, `register_pwmcc!`, the new
  `_register_sawtooth_tightening!`, and McCormick variants.
- Drop the redundant `MinMax((min = a, max = b))` wrapper — `MinMax` is a
  named-tuple alias.
- Vectorize the JuMP-side loops: `pwmcc_cuts` `brk`, sawtooth MIP
  constraints (4 vectorized `@constraint` families instead of a triple
  loop), epigraph LP and tangent constraints, manual-SOS2 adjacency, and
  the McCormick envelopes themselves.
- Extract `scale_back_g_basis(...)` (in `epigraph.jl`) and use it from
  both sawtooth and the epigraph tangent cuts: the
  `x_min² + (2·x_min·δ + δ²)·g₀ − Σ δ²·2^{−2k}·g_k` "scale back to actual
  dimensions" expression is shared.
- Restore the precomputed-form `add_bilinear_approx!` entrypoints for
  every bilinear method (Bin2, HybS, NMDT, DNMDT, NoBilinear) that takes
  pre-built x²/y² (or pre-built NMDT discretizations) instead of
  re-computing them. Backed by a private `_PrebuiltQuadApprox` adapter in
  `common.jl` for the bilinear methods that share math via a single
  `_build_*_with_precomputed` flow.
- Fix `sawtooth.jl`'s adjacency container schema (the old
  `(name, k, t)` container dropped the j axis when populated from a
  `(name, j, k, t)` loop). Now `(name, alpha_levels, 1:4, time)`, dense.
- NMDT binary-continuous-product McCormicks now register under a
  `_bc`-suffixed meta to avoid colliding with the residual product's
  McCormick under the same NMDT key.

All 1098 IOM tests pass.
The earlier consolidation commit on this branch edited signatures and
docstrings of `add_sparse_pwl_interpolation_variables!` and
`_add_generic_incremental_interpolation_constraint!` while moving the
file into src/approximations/. The path move was correct (everything
PWL-related belongs in approximations/), but the content edits were out
of scope — incremental.jl is a container-coupled PWL utility for HVDC
models, not an approximation method in the build/register sense.

Reverts the file content to byte-for-byte match main's
src/quadratic_approximations/incremental.jl while keeping the
consolidated path.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Establishes the layered pattern that the rest of the approximations will
follow:

  Layer 1 — pure-JuMP scalar math:
    build_mccormick_envelope(model, x, y, z, x_min, x_max, y_min, y_max)
      -> (upper_1, upper_2, lower_1, lower_2)
    build_mccormick_upper(model, x, y, z, x_min, x_max, y_min, y_max)
      -> (upper_1, upper_2)
    build_reformulated_mccormick(model, x, y, zp1, zx, zy, x_min, x_max, y_min, y_max)
      -> (c1, c2, c3, c4)

  Layer 2 — IOM adapter (allocate, loop, write):
    add_mccormick_approx!         -> McCormickConstraint   [name, 1:4, t]
    add_mccormick_upper_approx!   -> McCormickUpperConstraint   [name, 1:2, t]
    add_reformulated_mccormick_approx! -> ReformulatedMcCormickConstraint [name, 1:4, t]

The `lower_bounds` toggle is gone from the new API — callers that want
upper-only call `build_mccormick_upper` / `add_mccormick_upper_approx!`
directly. One math function maps 1:1 to one container key; the adapter
just iterates the returned NamedTuple values and slots them into the
container's 1:N inner axis.

Legacy vectorized `build_mccormick_envelope`/`build_reformulated_mccormick`
and the `register_mccormick_envelope!`/`register_reformulated_mccormick!`
helpers are preserved alongside until NMDT/Bin2/HybS migrate to the new
API in their own commits. Final cleanup of the legacy entry points
happens in the sweep phase.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Apply the McCormick template to the trivial cases. Each file now has:
  Scalar:  build_quadratic_approx(::NoQuadApproxConfig, model, x, x_min, x_max)
           build_bilinear_approx(::NoBilinearApproxConfig, model, x, y, x_min, x_max, y_min, y_max)
  Adapter: add_quadratic_approx!(::NoQuadApproxConfig, container, C, x_var, x_bounds, meta)
           add_bilinear_approx!(::NoBilinearApproxConfig, container, C, x_var, y_var, x_bounds, y_bounds, meta)

Scalar form returns a flat NamedTuple `(; approximation)` holding the exact
QuadExpr for that cell. The adapter loops `(name, t)`, calls the scalar,
and writes one cell at a time into a `QuadraticExpression` /
`BilinearProductExpression` container.

Legacy vectorized build_* + register_in_container! kept alongside for the
generic add_*_approx! wrapper in common.jl; the precomputed-form 12-arg
add_bilinear_approx! (used as a swap-in for Bin2/HybS) also kept. All
removed in the sweep phase.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Epigraph and sawtooth both have internal recursion-depth axes the math
owns (g_levels = 0:depth, alpha_levels = 1:depth, lp/mip/tangent sub-axes).
Scalar build_quadratic_approx(::Config, model, x_scalar, x_min, x_max)
takes one cell, builds the per-cell g/alpha/z variables and inner-axis
constraint arrays itself, returns them in a flat NamedTuple. IOM adapter
loops (name, t), calls the scalar per cell, slots refs into the
appropriately-shaped containers.

Sawtooth optionally composes epigraph for LP tightening. When enabled,
the scalar sawtooth calls scalar epigraph in its per-cell math and embeds
the epigraph NamedTuple in its `tightening` field. The sawtooth adapter
allocates the epigraph containers under `meta * "_lb"` and writes them
inline alongside sawtooth's own outputs.

Adds `scale_back_g_basis_scalar` helper (1D g_var, no name/t indexing) to
share the parabola anchor + residual decomposition between the new scalar
sawtooth and scalar epigraph paths.

Legacy result structs (EpigraphQuadResult, SawtoothQuadResult,
SawtoothTightening) + vectorized build + register_in_container! preserved
alongside until the generic add_quadratic_approx! wrapper in common.jl
goes away in the sweep phase.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
PWMCC: scalar `build_pwmcc_concave_cuts(model, v, q_expr, v_min, v_max, K)`
returns a NamedTuple of the K-segment δ/v^d binaries+continuous and the
selector/linking/interval-bound/chord/tangent constraints for one cell.

Solver-SOS2 and Manual-SOS2: scalar `build_quadratic_approx(::Config, model,
x, x_min, x_max)` returns the per-cell λ basis, link/norm/sos (or
adjacency) constraints, expression sums, and the approximation AffExpr.
When `pwmcc_segments > 0` the scalar calls `build_pwmcc_concave_cuts`
inline and embeds its NamedTuple in the `pwmcc` field.

Each adapter (`add_quadratic_approx!(::Config, ...)`) allocates the SOS2
containers AND (conditionally) the PWMCC containers under `meta * "_pwmcc"`,
loops `(name, t)`, calls the scalar per cell, and slots refs into the
appropriately-shaped containers.

Legacy result structs (PWMCCResult, SOS2QuadResult, ManualSOS2QuadResult),
legacy vectorized builds, and `register_pwmcc!` / `register_in_container!`
preserved alongside until callers migrate; removed in sweep.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Removes every vectorized `build_*` (acting on DenseAxisArrays of variables
and `Vector{MinMax}` bounds), every `register_in_container!` method, every
intermediate result struct (NoQuadApproxResult, EpigraphQuadResult,
SawtoothQuadResult, SawtoothTightening, SOS2QuadResult, ManualSOS2QuadResult,
PWMCCResult, NMDTDiscretization, NMDTBinaryContinuousProduct,
NMDTResidualProduct, NMDTEpigraphTightening, NMDTQuadResult, DNMDTQuadResult,
NMDTBilinearResult, DNMDTBilinearResult, Bin2BilinearResult, HybSBilinearResult),
the abstract QuadraticApproxResult/BilinearApproxResult supertypes,
`_PrebuiltQuadApprox`, `get_approximation`, vectorized `build_normed_variable`,
and the generic `add_quadratic_approx!`/`add_bilinear_approx!` dispatch
wrappers in common.jl.

Each method now ships:
  - scalar `build_<method>(model, x_scalar, x_min, x_max, ...)` returning a
    flat NamedTuple (pure JuMP, only @Assert preconditions — IS removed)
  - `add_<method>_approx!(container, ::Type{C}, x_var, x_bounds, meta)` IOM
    adapter that allocates known-axis containers, loops `(name, t)`, calls
    the scalar, slots refs into containers

Composed methods (bin2, hybs) call into the quad method's adapter for x²,
y², (x+y)², (x±y)², then do their own per-cell assembly. Each maintains a
precomputed-form `add_bilinear_approx!` that accepts already-built xsq/ysq
containers.

Test signatures updated to the new 6-arg
`add_quadratic_approx!(config, container, C, x_var, x_bounds, meta)` and
8-arg `add_bilinear_approx!(config, container, C, x_var, y_var, x_bounds,
y_bounds, meta)` shape — names and time_steps args dropped (axes come from
the var containers). test_pure_jump_approximations.jl removed.

All 1144 unit tests pass on HiGHS.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Comment on lines +54 to +57
model, x, y,
psq.approximation, xsq.approximation, ysq.approximation,
x_min, x_max, y_min, y_max,
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
model, x, y,
psq.approximation, xsq.approximation, ysq.approximation,
x_min, x_max, y_min, y_max,
)
model, x, y,
psq.approximation, xsq.approximation, ysq.approximation,
x_min, x_max, y_min, y_max,
)

Comment on lines +199 to +201
container, ReformulatedMcCormickConstraint, C,
name_axis, 1:4, time_axis; meta,
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
container, ReformulatedMcCormickConstraint, C,
name_axis, 1:4, time_axis; meta,
)
container, ReformulatedMcCormickConstraint, C,
name_axis, 1:4, time_axis; meta,
)

Comment on lines +100 to +101
model, x, y, z_var, x_min, x_max, y_min, y_max,
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
model, x, y, z_var, x_min, x_max, y_min, y_max,
)
model, x, y, z_var, x_min, x_max, y_min, y_max,
)

Comment on lines +234 to +235
container, McCormickConstraint, C, name_axis, 1:4, time_axis; meta,
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
container, McCormickConstraint, C, name_axis, 1:4, time_axis; meta,
)
container, McCormickConstraint, C, name_axis, 1:4, time_axis; meta,
)

Comment on lines +275 to +277
container, McCormickUpperConstraint, C,
name_axis, 1:depth, 1:2, time_axis; meta = mc_meta * "_lb",
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
container, McCormickUpperConstraint, C,
name_axis, 1:depth, 1:2, time_axis; meta = mc_meta * "_lb",
)
container, McCormickUpperConstraint, C,
name_axis, 1:depth, 1:2, time_axis; meta = mc_meta * "_lb",
)

Comment on lines +216 to +217
container, C, name_axis, time_axis, config.epigraph_depth, meta,
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
container, C, name_axis, time_axis, config.epigraph_depth, meta,
)
container, C, name_axis, time_axis, config.epigraph_depth, meta,
)

Comment on lines +288 to +289
container, C, name_axis, time_axis, config.epigraph_depth, meta,
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
container, C, name_axis, time_axis, config.epigraph_depth, meta,
)
container, C, name_axis, time_axis, config.epigraph_depth, meta,
)

@acostarelli
Copy link
Copy Markdown
Member Author

Closing because probably out of scope.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants