Pre-JOSS code review: correctness, robustness, and test fixes

[thevolatilebit]

Summary

A pre-submission code review surfaced ~15 severe and ~12 minor issues across the package. This PR lands the fixes for all 15 severe findings plus the applicable minors, with new unit tests covering previously untested public APIs (EfficientValueMDP, ConditionalUncertaintyReductionMDP, plot_ensemble_pareto, ensemble_to_dataframe, front).

The JOSS manuscript work is intentionally excluded from this PR.

Severe — correctness / crashes

• evaluate_experiments UndefVarError (StaticDesigns.jl:105). nrow(data) referenced an arg named X. Crashed on the default evaluate_empty_subset=true path whenever zero_cost_features was empty (also default). Tests never exercised it. Fixed; test added.
• SquaredMahalanobisDistance mixed axes (distancebased.jl). Λ_marginal was indexed in non_targets (column) order while vec_evidence and the row[evidence_keys] slice were in evidence-insertion order. Identical inputs produced different distances depending on insertion order (verified: 4.62 vs 5.90 with a non-trivial Σ). Now canonicalized against the column order at every cache key and dot product.
• front() dropped Pareto-optimal points and used a non-strict-weak sort (fronts.jl). The dominance check required strict improvement on both axes against the last kept point (rather than the partial order against all incumbents); with atol1=0 it silently lost points sharing an x-coordinate, and at non-zero atols it lost true front members below tolerance. The sort comparator used <= on ties → Base.sort undefined behavior. Replaced with a proper online Pareto-front algorithm; atols are now a separate post-hoc thinning step. Sort comparator fixed to <.
• EfficientValueMDP constructor checked the wrong value signature (EfficientValueMDP.jl:50). Asserted Tuple{Evidence, Vector{Float64}} while every call site passed NTuple{2,Float64}. Users who matched the docstring failed the assert; users who matched the assert crashed at runtime. Now checks Tuple{Evidence, NTuple{2, Float64}}.
• range(0.0, 1.0, thresholds) errored on thresholds=1 (UncertaintyReductionMDP.jl, ConditionalUncertaintyReductionMDP.jl). Now validates thresholds ≥ 2 with a clear ArgumentError.
• actions() filtered multi-feature experiments by first(features) only (all three MDPs). A multi-feature experiment was treated as "taken" iff only its first feature was observed; later observations could re-sample a feature already in evidence, or a fully observed experiment could remain available. Decided semantics: an experiment is treated as a single atomic action and excluded iff any of its features is already observed.
• Variance() / Entropy() divided by zero → silent NaN propagation (distancebased.jl). With a degenerate prior or constant target column, initial == 0 made the closure return NaN; the MDP terminal check uncertainty <= threshold was always false and the planner ran to max_experiments then took EOX with -bigM. Now validated at construction.
• All-similarities-zero recovery branch dropped active masks (distancebased.jl). The fallback reset similarities .= 1 and re-applied only the target hard-match, ignoring filter_range / importance_weights for evidence columns. Recovery now re-applies all active masks (or fails loudly when no rows match).
• EfficientValueMDP not exported (GenerativeDesigns.jl). The type had its own docstring and was the subject of efficient_value, but reviewers copying type-level examples hit UndefVarError. Added to export.
• zero_cost_features element-type assumption (StaticDesigns.jl). Splat into String[] failed for the common DataFrames Symbol idiom. Now accepts both.
• Generic catch in DistanceBased swallowed real errors (distancebased.jl). User-supplied distance closures had their stacktraces masked behind a single "unsupported scitype" message. The try is now narrowed to the dispatch on elscitype; user closures throw normally.
• target_condition numeric comparison broke on Multiclass / NaN (ConditionalUncertaintyReductionMDP.jl). Now dispatches on scitype and raises clearly on unsupported target types.
• sampler reproducibility hole (distancebased.jl). The default rng = default_rng() made repeated runs of efficient_design(s) / efficient_value non-deterministic with no public seed kwarg. The default is removed; rng is required.
• Mahalanobis exponential precomputation + silent singularity (distancebased.jl). Λ was eagerly computed for every non-empty subset (2^p − 1 inversions), and inv(Σ) returned Inf/NaN on rank-deficient submatrices without warning. Now lazy + memoized, with explicit singular detection (and a clear error suggesting diagonal > 0).

Minor

• Unit tests for previously untested public APIs: test/GenerativeDesigns/test_efficient_value.jl, test/GenerativeDesigns/test_conditional.jl, test/test_ensemble_fronts.jl, test/StaticDesigns/test.jl additions, expanded test/GenerativeDesigns/test_mahalanobis.jl, test/fronts.jl.
• Type stability of MDP structs: costs, uncertainty, weights, sampler, terminal_condition are now type parameters rather than abstract Function / untyped Dict, removing dynamic dispatch on every MCTS step.
• arrangements.jl guard against KeyError on partial evals when optimal_arrangement is run over a state outside the cached set.
• ensemble_to_dataframe action-set stringification: stopped relying on Julia's vector printing (["A"]["B"]); produces A,B directly.
• default_solver shared global: replaced with a default_solver() factory so MCTS RNG state isn't shared across calls.
• Inverted comment in test/GenerativeDesigns/test_active_sampling.jl fixed.
• Drop unused deps: Reexport, Requires removed from Project.toml.
• Empty-input guards in arrangements.jl and ensemble_fronts.jl.

Test plan

• Pkg.test() — 94 tests pass on Julia 1.12.5
• CI (Julia 1.10 / 1.12 matrix)
• Tutorials still render (ConditionalUncertaintyReduction, GliomaGrading, ActiveSampling)

[danielchen26]

The PR fixes several useful robustness issues, but there are still confirmed correctness, reproducibility, and API-consistency problems that can materially affect results. The most important blockers are: categorical distance direction, terminal_condition key normalization, RNG/solver reuse in generative paths, and lack of seeded reproducibility in the static design path. I would like to see these fixed with targeted regression tests before merging.

Thanks for the work in this PR. I reviewed the current diff in depth, including the distance kernels, generative MDPs, conditional MDP logic, static design path, ensemble front utilities, tests, docs, and the paper-facing behavior. Below are the issues I think should be addressed.

Blocking correctness / reproducibility issues

1. DiscreteDistance has the wrong direction

In src/GenerativeDesigns/distancebased.jl, the current implementation gives a larger distance to matching categorical values:

DiscreteDistance(; λ = 1) = function (x, col; _...)
    return map(y -> y == x ? λ : 0.0, col)
end

Because the default similarity is exponential decay in distance, matching categorical values receive lower posterior weight than non-matching values. That is opposite to the continuous distance convention, where an exact match gives distance 0.

This can directly corrupt default posterior weights for categorical features.

Suggested fix:

y == x ? 0.0 : λ

The docstring should also be updated from a match-based distance to a mismatch-based distance, for example λ * (x .!= col).

Please add a regression test showing that rows matching a categorical evidence value receive higher posterior weight than non-matching rows.

---

2. terminal_condition accepts Symbol keys at construction but fails at runtime

ConditionalUncertaintyReductionMDP validates condition columns with:

string(colname) in names(data)

but it stores the original terminal_condition unchanged. Later, conditional_likelihood checks the original key with:

colname in names(hist_data)

and indexes using that original key.

As a result, this can pass construction:

terminal_condition = (Dict(:Age => [30, 50]), 0.8)

but fail during isterminal.

Suggested fix:

Normalize terminal_condition during construction and store it as a concrete type, for example:

Dict{String, Tuple{Float64, Float64}}

The constructor should validate:

• the column exists;
• each range has exactly two values;
• both bounds are numeric;
• lower bound <= upper bound.

This would also address the current type instability from storing Tuple{Dict, Float64}.

Please add regression tests for Symbol keys and malformed ranges.

---

3. Generative MDP paths still reuse RNG and solver state

The PR improves the default solver by changing it from a shared constant to a fresh default_solver() call, which is a good direction. However, several generative paths still reuse mutable RNG and solver state.

Examples include:

queue = [Sim(mdp, planner; rng) for _ in 1:repetitions]

This passes the same RNG object into multiple simulations. These simulations are then executed through run_parallel, so this is a real source of nondeterminism / RNG state contention, not just a theoretical concern.

There is also solver reuse across threshold and ensemble loops, so ensemble runs are not fully independent.

Suggested fix:

• Derive one independent RNG per simulation before run_parallel.
• Prefer generating seeds up front, then using one RNG per Sim.
• Consider adding a backward-compatible solver_factory mechanism so each threshold / ensemble run can receive a fresh solver.
• Preserve compatibility for users who currently pass solver = DPWSolver(...).

The current tests do not cover this properly because many of them use N = 1 and/or repetitions = 0, which bypasses the risky paths.

Please add seeded reproducibility tests with N > 1 and repetitions > 1.

---

4. Static design path also lacks seeded reproducibility

The static design path also has a reproducibility issue. StaticDesigns.efficient_designs runs work under Threads.@threads, and optimal_arrangement constructs:

DPWSolver(; mdp_kwargs...)

without a seeded RNG entry point.

This means the static path can rely on global/default RNG state and become non-reproducible under threading.

Suggested fix:

• Add an rng or seed-based keyword.
• Pre-generate per-task seeds before entering threaded loops.
• Pass rng = Xoshiro(seed) into each DPWSolver.

Please add a seeded reproducibility test for the static design path.

---

Important robustness / semantic issues

5. QuadraticDistance does not handle zero-variance columns clearly

QuadraticDistance computes and caches σ:

σ = standardize ? var(col, Weights(prior); corrected = false) : 1
return λ * (x .- col) .^ 2 / σ

For constant columns, σ == 0, which can produce NaN or Inf. Depending on the data and readout, this can later fail with a low-level StatsBase error or produce misleading fallback behavior.

There is also a concern that σ is cached inside the closure. If the same distance closure is reused across multiple columns or priors, the first column’s variance may be reused incorrectly.

Suggested fix:

• If standardize && σ <= 0, throw a clear ArgumentError naming the offending column.
• Avoid caching σ across unrelated columns / prior inputs unless the closure is guaranteed to be column-specific.
• Add regression tests for constant predictor columns.

---

6. max_experiments counts evidence keys, not completed experiments

The baseline and conditional MDP action logic uses:

length(state.evidence) < m.max_experiments

This counts evidence entries, not completed experiments. That means:

• initial evidence can consume experiment budget;
• one multi-feature experiment can consume multiple slots;
• an experiment like ECG with several readout features can prematurely exhaust the budget.

This makes max_experiments semantically different from its name and documentation.

Suggested fix:

Count completed experiments instead of evidence keys. For example, count an experiment as completed when all of its associated features are present in evidence.

Please add tests covering:

• initial evidence;
• multi-feature experiments;
• max_experiments = 1 and max_experiments = 2.

---

7. realized_uncertainty=true mixes front-axis semantics and can inflate ensemble frequencies

When realized_uncertainty = true, the already-terminal branch of efficient_design / conditional_efficient_design returns a point whose second coordinate is realized uncertainty:

(0.0, realized_uncertainty)

but non-terminal designs use the nominal threshold:

(cost, threshold)

This means the front’s second coordinate mixes two different quantities.

Multiple thresholds above the initial uncertainty can collapse to the same (0.0, realized_uncertainty) point. front() itself is fine — retaining exact duplicate non-dominated points is mathematically valid — but downstream aggregation in ensemble_to_dataframe can then group these collapsed terminal designs and inflate frequency counts used for MLASP.

Suggested fix:

• Keep the front’s second coordinate consistently as the nominal threshold.
• Store realized uncertainty in metadata instead.
• Add a regression test showing that terminal designs do not inflate ensemble / MLASP frequency counts.

This is important for paper-facing ensemble / MLASP outputs.

---

8. Conditional constraints reject or mishandle missing values

The constructor currently rejects columns with eltype like:

Union{Missing, Float64}

because it checks:

eltype(col) <: Real

Real-world DataFrames often contain missing values. The later mask logic also does not explicitly coalesce missing to false.

Suggested fix:

• Use nonmissingtype(eltype(col)) <: Real.
• In the likelihood mask, treat missing as non-matching, for example with coalesce.(..., false).
• Document that rows with missing in constrained columns are excluded from the conditional likelihood.

Please add a regression test for numeric columns containing missing.

---

9. Public API paths should not use @assert for user input validation

Several public-facing paths use @assert for user input validation, including:

• conditional_likelihood;
• MDP constructor hasmethod checks;
• Entropy() scitype validation.

These should throw stable, user-facing ArgumentErrors instead of AssertionErrors.

Suggested fix:

Replace public-path assertions with explicit ArgumentErrors and update tests accordingly.

---

Lower-priority cleanup issues

10. importance_weights vector length is not validated

If a user passes an importance_weights vector with length different from nrow(data), the failure will occur later as a lower-level broadcasting or weighting error.

Suggested fix:

Validate the length during construction and throw a clear ArgumentError.

---

11. EfficientValueMDP documentation has the wrong value signature

The docs describe value as if it were a single-argument function of evidence, but the implementation requires:

value(evidence, costs)

Suggested fix:

Update the docstring and examples to consistently describe value(evidence, costs).

---

12. Documentation / defaults / naming are inconsistent

There are several smaller documentation and naming inconsistencies:

• Exponential docstring says λ = 1, but the code defaults to 1 / 2.
• terminal_condition defaults appear inconsistent across APIs / docs.
• thred_set appears to be a typo and should probably be renamed or aliased to tau_set / threshold_set.
• Typos such as acrss and Argumets should be fixed.

Suggested fix:

Clean up the docs and, where needed, add backwards-compatible aliases rather than breaking existing users.

---

13. Static feature normalization and plotting cleanup

zero_cost_features are normalized with string.(), but experiment feature names are not normalized the same way. This can still fail if experiment features are passed as Symbols.

Also:

• plot_front should guard against empty input.
• Average_Utility in ensemble output appears to represent cost / objective value rather than utility, so it should either be renamed or clearly documented.

Suggested fix:

• Normalize experiment feature names with string.().
• Add an empty-input guard to plot_front.
• Rename or document Average_Utility.

---

Paper / implementation notes

These are not blockers, but they should be documented clearly for reproducibility and JOSS-facing readers.

First, the implementation appears to normalize uncertainty by prior uncertainty. That explains why thresholds are swept over [0, 1]. This seems like a reasonable adaptation, but it should be documented explicitly.

Second, the paper’s terminal objective terms appear to be represented in the code as hard terminal constraints rather than as a literal terminal reward expression. That can be valid, but the distinction should be made clear for reproducibility.

---

Test coverage concern

The largest structural issue is that many tests are smoke / shape tests rather than correctness tests. They often verify that outputs are tuples or have fields, but they do not verify that the algorithms produce correct weights, rewards, terminal behavior, RNG reproducibility, or experiment-budget behavior.

Before merging, I would like to see targeted regression tests for at least:

1. categorical distance direction;
2. terminal_condition Symbol-key normalization;
3. missing values in conditional constraints;
4. constant-column QuadraticDistance;
5. seeded reproducibility for generative repetitions / ensembles;
6. seeded reproducibility for static threaded designs;
7. max_experiments with initial evidence and multi-feature experiments;
8. realized_uncertainty=true not inflating MLASP / ensemble frequency counts.

---

The PR fixes useful things, but the current version can still produce incorrect categorical posterior weights and non-reproducible results.

[danielchen26]

I have found above issues, can @thevolatilebit you check if those make sense, I double checked locally, they seems to be real issues. I could also open another branch and fix them and merge back here. Let me know.

[danielchen26]

@thevolatilebit I fixed the issues mentioned above. Please check

[thevolatilebit]

@danielchen26 Thank you! We can merge.

[slwu89]

Thanks both, looks like this solves some serious issues.

Just one comment, I noticed it adds a lot of code (+1,624 -255). We may want to try to keep PR LoC additions as small as possible to avoid saturating human or LLM context. But in this case it seems warranted.

Also, looks like we can delete the stale branch.