refactor: add missing parts of AbstractPopMember interface

Author: MilesCranmer

src/ExpressionBuilder.jl

@@ -11,7 +11,8 @@ using DynamicExpressions:
 using ..CoreModule: AbstractOptions, Dataset
 using ..HallOfFameModule: HallOfFame
 using ..PopulationModule: Population
-using ..PopMemberModule: PopMember
+using ..PopMemberModule: PopMember, AbstractPopMember, create_child
+using ..ComplexityModule: compute_complexity
 
 import DynamicExpressions: get_operators
 import ..CoreModule: create_expression
@@ -107,15 +108,16 @@ end
         return with_metadata(ex; init_params(options, dataset, ex, Val(true))...)
     end
     function embed_metadata(
-        member::PopMember, options::AbstractOptions, dataset::Dataset{T,L}
-    ) where {T,L}
-        return PopMember(
+        member::PM, options::AbstractOptions, dataset::Dataset{T,L}
+    ) where {T,L,N,PM<:AbstractPopMember{T,L,N}}
+        return create_child(
+            member,
             embed_metadata(member.tree, options, dataset),
             member.cost,
             member.loss,
-            nothing;
-            member.ref,
-            member.parent,
+            options;
+            complexity=compute_complexity(member, options),
+            parent_ref=member.ref,
             deterministic=options.deterministic,
         )
     end
@@ -135,7 +137,7 @@ end
     end
     function embed_metadata(
         vec::Vector{H}, options::AbstractOptions, dataset::Dataset{T,L}
-    ) where {T,L,H<:Union{HallOfFame,Population,PopMember}}
+    ) where {T,L,H<:Union{HallOfFame,Population,AbstractPopMember}}
         return map(Fix{2}(Fix{3}(embed_metadata, dataset), options), vec)
     end
 end

src/HallOfFame.jl

@@ -73,17 +73,36 @@ function HallOfFame(
     options::AbstractOptions, dataset::Dataset{T,L}
 ) where {T<:DATA_TYPE,L<:LOSS_TYPE}
     base_tree = create_expression(init_value(T), options, dataset)
+    PM = options.popmember_type
 
-    return HallOfFame{T,L,typeof(base_tree),PopMember{T,L,typeof(base_tree)}}(
+    # Create a prototype member to get the concrete type
+    prototype = PM(
+        copy(base_tree),
+        L(0),
+        L(Inf),
+        options,
+        1;  # complexity
+        parent=-1,
+        deterministic=options.deterministic,
+    )
+
+    PMtype = typeof(prototype)
+
+    return HallOfFame{T,L,typeof(base_tree),PMtype}(
         [
-            PopMember(
-                copy(base_tree),
-                L(0),
-                L(Inf),
-                options;
-                parent=-1,
-                deterministic=options.deterministic,
-            ) for i in 1:(options.maxsize)
+            if i == 1
+                prototype
+            else
+                PM(
+                    copy(base_tree),
+                    L(0),
+                    L(Inf),
+                    options,
+                    1;  # complexity
+                    parent=-1,
+                    deterministic=options.deterministic,
+                )
+            end for i in 1:(options.maxsize)
         ],
         [false for i in 1:(options.maxsize)],
     )

src/Mutate.jl

@@ -40,10 +40,10 @@ using ..MutationFunctionsModule:
 using ..ConstantOptimizationModule: optimize_constants
 using ..RecorderModule: @recorder
 
-abstract type AbstractMutationResult{N<:AbstractExpression,P<:PopMember} end
+abstract type AbstractMutationResult{N<:AbstractExpression,P<:AbstractPopMember} end
 
 """
-    MutationResult{N<:AbstractExpression,P<:PopMember}
+    MutationResult{N<:AbstractExpression,P<:AbstractPopMember}
 
 Represents the result of a mutation operation in the genetic programming algorithm. This struct is used to return values from `mutate!` functions.
 
@@ -160,7 +160,7 @@ Use this to modify how `mutate_constant` changes for an expression type.
 function condition_mutate_constant!(
     ::Type{<:AbstractExpression},
     weights::AbstractMutationWeights,
-    member::PopMember,
+    member::AbstractPopMember,
     options::AbstractOptions,
     curmaxsize::Int,
 )
@@ -352,7 +352,7 @@ end
 
 @generated function _dispatch_mutations!(
     tree::AbstractExpression,
-    member::PopMember,
+    member::AbstractPopMember,
     mutation_choice::Symbol,
     weights::W,
     options::AbstractOptions;
@@ -381,7 +381,7 @@ end
         mutation_weights::AbstractMutationWeights,
         options::AbstractOptions;
         kws...,
-    ) where {N<:AbstractExpression,P<:PopMember,S}
+    ) where {N<:AbstractExpression,P<:AbstractPopMember,S}
 
 Perform a mutation on the given `tree` and `member` using the specified mutation type `S`.
 Various `kws` are provided to access other data needed for some mutations.
@@ -409,7 +409,7 @@ so it can always return immediately.
 """
 function mutate!(
     ::N, ::P, ::Val{S}, ::AbstractMutationWeights, ::AbstractOptions; kws...
-) where {N<:AbstractExpression,P<:PopMember,S}
+) where {N<:AbstractExpression,P<:AbstractPopMember,S}
     return error("Unknown mutation choice: $S")
 end
 
@@ -422,7 +422,7 @@ function mutate!(
     recorder::RecordType,
     temperature,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = mutate_constant(tree, temperature, options)
     @recorder recorder["type"] = "mutate_constant"
     return MutationResult{N,P}(; tree=tree)
@@ -436,7 +436,7 @@ function mutate!(
     options::AbstractOptions;
     recorder::RecordType,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = mutate_operator(tree, options)
     @recorder recorder["type"] = "mutate_operator"
     return MutationResult{N,P}(; tree=tree)
@@ -451,7 +451,7 @@ function mutate!(
     recorder::RecordType,
     nfeatures,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = mutate_feature(tree, nfeatures)
     @recorder recorder["type"] = "mutate_feature"
     return MutationResult{N,P}(; tree=tree)
@@ -465,7 +465,7 @@ function mutate!(
     options::AbstractOptions;
     recorder::RecordType,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = swap_operands(tree)
     @recorder recorder["type"] = "swap_operands"
     return MutationResult{N,P}(; tree=tree)
@@ -480,7 +480,7 @@ function mutate!(
     recorder::RecordType,
     nfeatures,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     if rand() < 0.5
         tree = append_random_op(tree, options, nfeatures)
         @recorder recorder["type"] = "add_node:append"
@@ -500,7 +500,7 @@ function mutate!(
     recorder::RecordType,
     nfeatures,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = insert_random_op(tree, options, nfeatures)
     @recorder recorder["type"] = "insert_node"
     return MutationResult{N,P}(; tree=tree)
@@ -514,7 +514,7 @@ function mutate!(
     options::AbstractOptions;
     recorder::RecordType,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = delete_random_op!(tree)
     @recorder recorder["type"] = "delete_node"
     return MutationResult{N,P}(; tree=tree)
@@ -528,7 +528,7 @@ function mutate!(
     options::AbstractOptions;
     recorder::RecordType,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = form_random_connection!(tree)
     @recorder recorder["type"] = "form_connection"
     return MutationResult{N,P}(; tree=tree)
@@ -542,7 +542,7 @@ function mutate!(
     options::AbstractOptions;
     recorder::RecordType,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = break_random_connection!(tree)
     @recorder recorder["type"] = "break_connection"
     return MutationResult{N,P}(; tree=tree)
@@ -556,7 +556,7 @@ function mutate!(
     options::AbstractOptions;
     recorder::RecordType,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     tree = randomly_rotate_tree!(tree)
     @recorder recorder["type"] = "rotate_tree"
     return MutationResult{N,P}(; tree=tree)
@@ -572,7 +572,7 @@ function mutate!(
     recorder::RecordType,
     parent_ref,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     @assert options.should_simplify
     simplify_tree!(tree, options.operators)
     tree = combine_operators(tree, options.operators)
@@ -593,7 +593,7 @@ function mutate!(
     curmaxsize,
     nfeatures,
     kws...,
-) where {T,N<:AbstractExpression{T},P<:PopMember}
+) where {T,N<:AbstractExpression{T},P<:AbstractPopMember}
     tree = randomize_tree(tree, curmaxsize, options, nfeatures)
     @recorder recorder["type"] = "randomize"
     return MutationResult{N,P}(; tree=tree)
@@ -608,7 +608,7 @@ function mutate!(
     recorder::RecordType,
     dataset::Dataset,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     cur_member, new_num_evals = optimize_constants(dataset, member, options)
     @recorder recorder["type"] = "optimize"
     return MutationResult{N,P}(;
@@ -625,7 +625,7 @@ function mutate!(
     recorder::RecordType,
     parent_ref,
     kws...,
-) where {N<:AbstractExpression,P<:PopMember}
+) where {N<:AbstractExpression,P<:AbstractPopMember}
     @recorder begin
         recorder["type"] = "identity"
         recorder["result"] = "accept"

src/ParametricExpression.jl

@@ -24,7 +24,7 @@ using ..CoreModule:
     AbstractExpressionSpec,
     get_indices,
     ExpressionSpecModule as ES
-using ..PopMemberModule: PopMember
+using ..PopMemberModule: PopMember, AbstractPopMember
 using ..InterfaceDynamicExpressionsModule: InterfaceDynamicExpressionsModule as IDE
 using ..LossFunctionsModule: LossFunctionsModule as LF
 using ..ExpressionBuilderModule: ExpressionBuilderModule as EB
@@ -102,7 +102,7 @@ end
 function MM.condition_mutate_constant!(
     ::Type{<:ParametricExpression},
     weights::AbstractMutationWeights,
-    member::PopMember,
+    member::AbstractPopMember,
     options::AbstractOptions,
     curmaxsize::Int,
 )

src/PopMember.jl

@@ -198,35 +198,25 @@ function recompute_complexity!(
     return complexity
 end
 
-# Interface for creating child members with custom field preservation
 """
-    create_child(parent::P, tree, cost, loss, options;
-                complexity::Union{Int,Nothing}=nothing, parent_ref, kwargs...) where P<:AbstractPopMember
-
-Create a new PopMember derived from a parent (mutation case).
-Custom types should override to preserve their additional fields.
+    create_child(parent::P, tree::AbstractExpression{T}, cost, loss, options;
+                complexity::Union{Int,Nothing}=nothing, parent_ref, kwargs...) where {T,L,P<:PopMember{T,L}}
 
-# Arguments
-- `parent`: The parent member being mutated
-- `tree`: The new expression tree
-- `cost`: The new cost
-- `loss`: The new loss
-- `options`: The options
-- `complexity`: The complexity (computed if not provided)
-- `parent_ref`: Reference to parent for tracking
+Create a new PopMember with a potentially different expression type.
+Used by embed_metadata where the expression gains metadata.
 """
 function create_child(
     parent::P,
-    tree::N,
+    tree::AbstractExpression{T},
     cost::L,
     loss::L,
     options;
     complexity::Union{Int,Nothing}=nothing,
     parent_ref,
     kwargs...,
-) where {T,L,N<:AbstractExpression{T},P<:PopMember{T,L,N}}
+) where {T,L,P<:PopMember{T,L}}
     actual_complexity = @something complexity compute_complexity(tree, options)
-    return constructorof(P)(
+    return PopMember(
         tree,
         cost,
         loss,
@@ -246,16 +236,16 @@ Custom types should override to blend their additional fields.
 """
 function create_child(
     parents::Tuple{P,P},
-    tree::N,
+    tree::AbstractExpression{T},
     cost::L,
     loss::L,
     options;
     complexity::Union{Int,Nothing}=nothing,
     parent_ref,
     kwargs...,
-) where {T,L,N<:AbstractExpression{T},P<:PopMember{T,L,N}}
+) where {T,L,P<:PopMember{T,L}}
     actual_complexity = @something complexity compute_complexity(tree, options)
-    return constructorof(P)(
+    return PopMember(
         tree,
         cost,
         loss,