Log all complexities over time

src/MLJInterface.jl

@@ -41,6 +41,8 @@ function modelexpr(model_name::Symbol)
  addprocs_function::Union{Function,Nothing} = nothing
  heap_size_hint_in_bytes::Union{Integer,Nothing} = nothing
  logger::Union{AbstractLogger,Nothing} = nothing
+ logging_callback::Union{Function,Nothing} = nothing
+ log_every_n::Int = 1
  runtests::Bool = true
  loss_type::L = Nothing
  selection_method::Function = choose_best
@@ -177,6 +179,8 @@ function _update(m, verbosity, old_fitresult, old_cache, X, y, w, options)
  y_units=y_units_clean,
  verbosity=verbosity,
  logger=m.logger,
+ logging_callback=m.logging_callback,
+ log_every_n=m.log_every_n,
  # Help out with inference:
  v_dim_out=isa(m, SRRegressor) ? Val(1) : Val(2),
  )

src/SearchUtils.jl

@@ -407,7 +407,7 @@ end
 function default_logging_callback(logger; options, num_evals, hall_of_fame, datasets, _...)
  L = typeof(first(datasets).baseline_loss)
  with_logger(logger) do
- @info("search_state", num_evals = sum(sum, num_evals))
+ d = Dict()
  for (i, (hof, dataset)) in enumerate(zip(hall_of_fame, datasets))
  dominating = calculate_pareto_frontier(hof)
  best_loss = length(dominating) > 0 ? dominating[end].loss : L(Inf)
@@ -417,15 +417,17 @@ function default_logging_callback(logger; options, num_evals, hall_of_fame, data
  string_tree(member.tree, options; variable_names=dataset.variable_names) for
  member in dominating
  ]
- @info(
-  "search_state_$(i)",
-  best_loss = best_loss,
-  equations = equations,
- losses = losses,
- complexities = complexities,
- log_step_increment = 0,
- )
+ d[string(i)] = Dict()
+ d[string(i)]["best_loss"] = best_loss
+ d[string(i)]["equations"] = Dict()
+ for (complexity, loss, equation) in zip(complexities, losses, equations)
+ d[string(i)]["equations"][string(complexity)] = Dict(
+  "loss" => loss, "equation" => equation
+ )
+ end
  end
+ d["num_evals"] = sum(sum, num_evals)
+ @info("search_state", data = d)
  end
 end
 

src/SymbolicRegression.jl

@@ -355,6 +355,7 @@ function equation_search(
  verbosity::Union{Integer,Nothing}=nothing,
  logger::Union{AbstractLogger,Nothing}=nothing,
  logging_callback::Union{Function,Nothing}=nothing,
+ log_every_n::Int=1,
  progress::Union{Bool,Nothing}=nothing,
  X_units::Union{AbstractVector,Nothing}=nothing,
  y_units=nothing,
@@ -403,6 +404,7 @@ function equation_search(
  verbosity=verbosity,
  logger=logger,
  logging_callback=logging_callback,
+ log_every_n=log_every_n,
  progress=progress,
  v_dim_out=Val(DIM_OUT),
  )
@@ -574,6 +576,7 @@ function equation_search(
  saved_state,
  _verbosity,
  _logging_callback,
+ log_every_n,
  _progress,
  Val(_return_state),
  )
@@ -593,6 +596,7 @@ function _equation_search(
  saved_state,
  verbosity,
  logging_callback,
+ log_every_n,
  progress,
  ::Val{RETURN_STATE},
 ) where {T<:DATA_TYPE,L<:LOSS_TYPE,D<:Dataset{T,L},PARALLELISM,RETURN_STATE,DIM_OUT}
@@ -800,6 +804,7 @@ function _equation_search(
  )
  end
 
+ log_step = 0
  last_print_time = time()
  last_speed_recording_time = time()
  num_evals_last = sum(sum, num_evals)
@@ -971,7 +976,7 @@ function _equation_search(
  PARALLELISM,
  )
  end
- if logging_callback !== nothing
+ if logging_callback !== nothing && log_step % log_every_n == 0
  logging_callback(;
  options,
  num_evals,
@@ -982,6 +987,7 @@ function _equation_search(
  datasets=datasets,
  )
  end
+ log_step += 1
  end
  sleep(1e-6)