Merge pull request #494 from MilesCranmer/fix-mlj-option-caching

fix: caching of options in MLJ regressors

Author: MilesCranmer

src/Core.jl

@@ -30,7 +30,9 @@ using .OptionsStructModule:
     Options,
     ComplexityMapping,
     specialized_options,
-    operator_specialization
+    operator_specialization,
+    WarmStartIncompatibleError,
+    check_warm_start_compatibility
 using .OperatorsModule:
     get_safe_op,
     plus,

src/MLJInterface.jl

@@ -37,7 +37,8 @@ using ..CoreModule:
     ComplexityMapping,
     AbstractExpressionSpec,
     ExpressionSpec,
-    get_expression_type
+    get_expression_type,
+    check_warm_start_compatibility
 using ..CoreModule.OptionsModule: DEFAULT_OPTIONS, OPTION_DESCRIPTIONS
 using ..ComplexityModule: compute_complexity
 using ..HallOfFameModule: HallOfFame, format_hall_of_fame
@@ -232,7 +233,10 @@ function MMI.update(
     y,
     w=nothing,
 )
-    options = old_fitresult === nothing ? get_options(m) : old_fitresult.options
+    options = get_options(m)
+    if !isnothing(old_fitresult)
+        check_warm_start_compatibility(old_fitresult.options, options)
+    end
     return _update(m, verbosity, old_fitresult, old_cache, X, y, w, options, nothing)
 end
 function _update(

src/Options.jl

@@ -1038,6 +1038,7 @@ $(OPTION_DESCRIPTIONS)
     }(
         operators,
         op_constraints,
+        _nested_constraints,
         complexity_mapping,
         tournament_selection_n,
         tournament_selection_p,
@@ -1099,7 +1100,6 @@ $(OPTION_DESCRIPTIONS)
         max_evals,
         input_stream,
         skip_mutation_failures,
-        _nested_constraints,
         deterministic,
         define_helper_functions,
         use_recorder,

src/OptionsStruct.jl

@@ -28,6 +28,13 @@ end
 
 Base.eltype(::ComplexityMapping{T}) where {T} = T
 
+function Base.:(==)(a::ComplexityMapping, b::ComplexityMapping)
+    return a.use == b.use &&
+           a.op_complexities == b.op_complexities &&
+           a.variable_complexity == b.variable_complexity &&
+           a.constant_complexity == b.constant_complexity
+end
+
 """Promote type when defining complexity mapping."""
 function ComplexityMapping(;
     op_complexities::Tuple{Vararg{Vector,D}},
@@ -182,6 +189,7 @@ struct Options{
 } <: AbstractOptions
     operators::OP
     op_constraints::OP_CONSTRAINTS
+    nested_constraints::Union{Vector{Tuple{Int,Int,Vector{Tuple{Int,Int,Int}}}},Nothing}
     complexity_mapping::CM
     tournament_selection_n::Int
     tournament_selection_p::Float32
@@ -243,7 +251,6 @@ struct Options{
     max_evals::Union{Int,Nothing}
     input_stream::IO
     skip_mutation_failures::Bool
-    nested_constraints::Union{Vector{Tuple{Int,Int,Vector{Tuple{Int,Int,Int}}}},Nothing}
     deterministic::Bool
     define_helper_functions::Bool
     use_recorder::Bool
@@ -290,4 +297,40 @@ end
     end
 end
 
+struct WarmStartIncompatibleError <: Exception
+    fields::Vector{Symbol}
+end
+
+function Base.showerror(io::IO, e::WarmStartIncompatibleError)
+    print(io, "Warm start incompatible due to changed field(s): ")
+    join(io, e.fields, ", ")
+    return print(io, ". Use `fit!(mach, force=true)` to restart training.")
+end
+
+check_warm_start_compatibility(::AbstractOptions, ::AbstractOptions) = nothing  # LCOV_EXCL_LINE
+
+function check_warm_start_compatibility(old_options::Options, new_options::Options)
+    incompatible_fields = (
+        :operators,
+        :op_constraints,
+        :nested_constraints,
+        :complexity_mapping,
+        :dimensionless_constants_only,
+        :maxsize,
+        :maxdepth,
+        :populations,
+        :population_size,
+        :node_type,
+        :expression_type,
+        :expression_options,
+    )
+
+    changed = [
+        f for f in incompatible_fields if
+        getproperty(old_options, f) != getproperty(new_options, f)
+    ]
+    isempty(changed) || throw(WarmStartIncompatibleError(changed))
+    return nothing
+end
+
 end

src/SymbolicRegression.jl

@@ -243,6 +243,7 @@ using .CoreModule:
     AbstractOptions,
     Options,
     ComplexityMapping,
+    WarmStartIncompatibleError,
     AbstractMutationWeights,
     MutationWeights,
     AbstractExpressionSpec,

test/test_mlj.jl

@@ -257,3 +257,36 @@ end
     end
     @test occursin("Evaluation failed either due to", msg)
 end
+
+@testitem "MLJ options caching fix" tags = [:part3] begin
+    using SymbolicRegression
+    using SymbolicRegression: WarmStartIncompatibleError
+    using MLJBase
+    using Random: MersenneTwister
+    using Suppressor
+
+    include("test_params.jl")
+
+    # Test that parameter changes are respected and incompatible changes throw errors
+    rng = MersenneTwister(0)
+    X = (x1=randn(rng, 50), x2=randn(rng, 50))
+    y = @. 2.0 * X.x1 + 3.0 * X.x2
+
+    model = SRRegressor(;
+        binary_operators=[+, -, *], niterations=2, tournament_selection_n=10, populations=2
+    )
+
+    mach = machine(model, X, y)
+    @suppress fit!(mach, verbosity=0)
+
+    # Test compatible parameter change
+    model.tournament_selection_n = 20
+    @suppress fit!(mach, verbosity=0)
+    @test mach.fitresult.options.tournament_selection_n == 20  # Should be updated
+
+    # Test incompatible parameter change throws error with correct message
+    model.populations = 4
+    err = @test_throws WarmStartIncompatibleError @suppress fit!(mach, verbosity=0)
+    @test :populations ∈ err.value.fields
+    @test occursin("force=true", sprint(showerror, err.value))
+end