refactor: infer_popmember_type

Author: MilesCranmer

src/PopMember.jl

@@ -2,7 +2,7 @@ module PopMemberModule
 
 using DispatchDoctor: @unstable
 using DynamicExpressions: AbstractExpression, AbstractExpressionNode, string_tree
-import DynamicExpressions: constructorof
+import DynamicExpressions: constructorof, with_type_parameters
 using ..CoreModule: AbstractOptions, Dataset, DATA_TYPE, LOSS_TYPE, create_expression
 import ..CoreModule.OptionsModule: default_popmember_type
 import ..ComplexityModule: compute_complexity
@@ -267,4 +267,14 @@ function popmember_type end
     return PopMember{T,L,N}
 end
 
+@inline function with_type_parameters(
+    ::Type{<:PopMember}, ::Type{T}, ::Type{L}, ::Type{N}
+) where {T,L,N}
+    return PopMember{T,L,N}
+end
+
+@inline function expression_type(::Type{<:AbstractPopMember{<:Any,<:Any,N}}) where {N}
+    return N
+end
+
 end

src/SearchUtils.jl

@@ -23,7 +23,7 @@ using ..CoreModule:
     Dataset, AbstractOptions, Options, RecordType, max_features, create_expression
 using ..ComplexityModule: compute_complexity
 using ..PopulationModule: Population
-using ..PopMemberModule: PopMember, AbstractPopMember, with_expression_type
+using ..PopMemberModule: PopMember, AbstractPopMember
 using ..HallOfFameModule: HallOfFame, string_dominating_pareto_curve
 using ..ConstantOptimizationModule: optimize_constants
 using ..ProgressBarsModule: WrappedProgressBar, manually_iterate!, barlen
@@ -34,6 +34,15 @@ using ..CheckConstraintsModule: check_constraints
 
 function logging_callback! end
 
+@unstable @inline function infer_popmember_type(
+    ::Type{T}, ::Type{L}, ::Type{D}, options
+) where {T,L,D<:Dataset}
+    NodeType = with_type_parameters(options.node_type, T)
+    N = Base.promote_op(create_expression, NodeType, typeof(options), D)
+    N in (Any, Union{}) && error("Failed to infer expression type")
+    return with_type_parameters(options.popmember_type, T, L, N)
+end
+
 """
     @filtered_async expr
 
@@ -793,26 +802,13 @@ end
 
 """Parse user-provided guess expressions and convert them into optimized
 `PopMember` objects for each output dataset."""
-@unstable function parse_guesses(
-    ::Type{P},
-    guesses::Union{AbstractVector,AbstractVector{<:AbstractVector}},
-    datasets::Vector{D},
-    options::AbstractOptions,
-) where {T,L,N,P<:AbstractPopMember{T,L,N},D<:Dataset{T,L}}
-    return _parse_guesses_impl(P, guesses, datasets, options)
-end
-
-# Deal with non-concrete PopMember types
 @unstable function parse_guesses(
     ::Type{P},
     guesses::Union{AbstractVector,AbstractVector{<:AbstractVector}},
     datasets::Vector{D},
     options::AbstractOptions,
 ) where {T,L,P<:AbstractPopMember{T,L},D<:Dataset{T,L}}
-    NodeType = with_type_parameters(options.node_type, T)
-    N = Base.promote_op(create_expression, NodeType, typeof(options), D)
-    N in (Any, Union{}) && error("Failed to infer expression type")
-    ConcreteP = with_expression_type(P, N)
+    ConcreteP = infer_popmember_type(T, L, D, options)
     return _parse_guesses_impl(ConcreteP, guesses, datasets, options)
 end
 

src/SymbolicRegression.jl

@@ -297,7 +297,8 @@ using .MutationFunctionsModule:
 using .InterfaceDynamicExpressionsModule:
     @extend_operators, require_copy_to_workers, make_example_inputs
 using .LossFunctionsModule: eval_loss, eval_cost, update_baseline_loss!, score_func
-using .PopMemberModule: AbstractPopMember, PopMember, reset_birth!, popmember_type
+using .PopMemberModule:
+    AbstractPopMember, PopMember, reset_birth!, popmember_type, expression_type
 using .CoreModule.UtilsModule: get_birth_order
 using .PopulationModule: Population, best_sub_pop, record_population, best_of_sample
 using .HallOfFameModule:
@@ -339,7 +340,8 @@ using .SearchUtilsModule:
     get_cur_maxsize,
     update_hall_of_fame!,
     parse_guesses,
-    logging_callback!
+    logging_callback!,
+    infer_popmember_type
 using .LoggingModule: AbstractSRLogger, SRLogger, get_logger
 using .TemplateExpressionModule:
     TemplateExpression, TemplateStructure, TemplateExpressionSpec, ParamVector, has_params
@@ -631,20 +633,8 @@ end
     @recorder record["options"] = "$(options)"
 
     nout = length(datasets)
-    example_dataset = first(datasets)
-    example_ex = create_expression(init_value(T), options, example_dataset)
-    NT = typeof(example_ex)
-    # Create a prototype member to get the concrete type
-    prototype_member = options.popmember_type(
-        copy(example_ex),
-        L(0),
-        L(Inf),
-        options,
-        1;  # complexity
-        parent=-1,
-        deterministic=options.deterministic,
-    )
-    PMType = typeof(prototype_member)
+    PMType = infer_popmember_type(T, L, D, options)
+    NT = expression_type(PMType)
     PopType = Population{T,L,NT,PMType}
     HallOfFameType = HallOfFame{T,L,NT,PMType}
     WorkerOutputType = get_worker_output_type(

test/test_abstract_popmember.jl

@@ -79,6 +79,13 @@
 
     @unstable DynamicExpressions.constructorof(::Type{<:CustomPopMember}) = CustomPopMember
 
+    # Define with_type_parameters for CustomPopMember
+    @unstable function DynamicExpressions.with_type_parameters(
+        ::Type{<:CustomPopMember}, ::Type{T}, ::Type{L}, ::Type{N}
+    ) where {T,L,N}
+        return CustomPopMember{T,L,N}
+    end
+
     # Define copy for CustomPopMember
     function Base.copy(p::CustomPopMember)
         return CustomPopMember(