Merge pull request #226 from MilesCranmer/mlj-integration

MLJ Integration

Author: MilesCranmer

Project.toml

@@ -1,7 +1,7 @@
 name = "SymbolicRegression"
 uuid = "8254be44-1295-4e6a-a16d-46603ac705cb"
 authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
-version = "0.19.1"
+version = "0.20.0"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
@@ -10,6 +10,8 @@ DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 LossFunctions = "30fc2ffe-d236-52d8-8643-a9d8f7c094a7"
+MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
@@ -33,6 +35,8 @@ DynamicExpressions = "0.9"
 JSON3 = "1"
 LineSearches = "7"
 LossFunctions = "0.6, 0.7, 0.8, 0.10"
+MLJModelInterface = "1.5, 1.6, 1.7, 1.8"
+MacroTools = "0.4, 0.5"
 Optim = "0.19, 1.1"
 Pkg = "1"
 PrecompileTools = "1"
@@ -47,10 +51,12 @@ julia = "1.6"
 [extras]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d"
+MLJTestInterface = "72560011-54dd-4dc2-94f3-c5de45b75ecd"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "SafeTestsets", "ForwardDiff", "LinearAlgebra", "SymbolicUtils", "Zygote"]
+test = ["Test", "SafeTestsets", "ForwardDiff", "LinearAlgebra", "MLJBase", "MLJTestInterface", "SymbolicUtils", "Zygote"]

src/HallOfFame.jl

@@ -161,4 +161,44 @@ function string_dominating_pareto_curve(
     return output
 end
 
+function format_hall_of_fame(
+    hof::HallOfFame{T,L}, options, baseline_loss::L
+) where {T<:DATA_TYPE,L<:LOSS_TYPE}
+    ZERO_POINT = L(1e-10)
+
+    dominating = calculate_pareto_frontier(hof)
+    cur_loss = baseline_loss
+    last_loss = cur_loss
+    last_complexity = 0
+
+    trees = [member.tree for member in dominating]
+    losses = [member.loss for member in dominating]
+    complexities = [compute_complexity(member, options) for member in dominating]
+    scores = Array{L}(undef, length(dominating))
+
+    for i in 1:length(dominating)
+        complexity = complexities[i]
+        cur_loss = losses[i]
+        delta_c = complexity - last_complexity
+        delta_l_mse = log(abs(cur_loss / last_loss) + ZERO_POINT)
+
+        scores[i] = -delta_l_mse / delta_c
+        last_loss = cur_loss
+        last_complexity = complexity
+    end
+    return (; trees=trees, scores=scores, losses=losses, complexities=complexities)
+end
+function format_hall_of_fame(
+    hof::AH, options, baseline_loss
+) where {T,L,H<:HallOfFame{T,L},AH<:AbstractVector{H}}
+    outs = [format_hall_of_fame(h, options, baseline_loss) for h in hof]
+    return (;
+        trees=[out.trees for out in outs],
+        scores=[out.scores for out in outs],
+        losses=[out.losses for out in outs],
+        complexities=[out.complexities for out in outs],
+    )
+end
+# TODO: Re-use this in `string_dominating_pareto_curve`
+
 end