DataDrivenLux

AbstractDAGSRAlgorithm

Developer interface for differentiable directed-acyclic-graph symbolic-regression algorithms. This interface is intended for solver extensions, not ordinary users.

Interface

A subtype must provide an options field compatible with CommonAlgOptions and methods init_model(alg, basis, dataset, intervals) and update_parameters!(cache::SearchCache{<:MyAlgorithm}). The generic cache initialization supplies the dataset, candidate population, and optimization state. init_model must return a callable Lux model compatible with the basis and dataset dimensions. update_parameters! must mutate cache.p or other algorithm state in place and return nothing. An algorithm that uses the default layered graph can reuse the generic init_model method.

The init_model method should retain the package's dispatch shape, (::MyAlgorithm, ::Basis, ::Dataset, intervals), so it is more specific than the default method while remaining applicable to the common solver path.

The generic CommonSolve.solve! path consumes the cache, repeatedly calls update_parameters!, and returns a DataDrivenDiffEq.DataDrivenSolution. Custom algorithms should keep the loss, keep, and population semantics of CommonAlgOptions or document any intentional differences.

Example

struct MyDAGAlgorithm <: DataDrivenLux.AbstractDAGSRAlgorithm
    options::DataDrivenLux.CommonAlgOptions
end

DataDrivenLux.init_model(alg, basis, dataset, intervals) =
    DataDrivenLux.LayeredDAG(
        length(basis), size(dataset.y, 1), 1, (1,), (identity,)
    )
DataDrivenLux.update_parameters!(cache::DataDrivenLux.SearchCache{<:MyDAGAlgorithm}) = nothing

CommonAlgOptions(; kwargs...)

Shared configuration for AbstractDAGSRAlgorithm implementations. Concrete algorithms normally expose these keywords through their own constructor.

Fields

• populationsize

• functions

• arities

• n_layers

• skip

• simplex

• loss

• keep

• use_protected

• distributed

• threaded

• rng

• optimizer

• optim_options

• optimiser

• observed

• alpha

Keywords

• populationsize::Int: number of candidate graphs retained in the population.
• functions: candidate unary and binary functions.
• arities: arity corresponding to each entry in functions.
• n_layers::Int: number of learned graph layers.
• skip::Bool: whether each layer receives skip connections.
• simplex::AbstractSimplex: map used for categorical path weights.
• loss: function used to rank candidates.
• keep::Union{Real,Int}: retained fraction or number of candidates.
• use_protected::Bool: whether unsafe symbolic operations are replaced by safe versions.
• distributed::Bool: whether candidate optimization uses worker processes.
• threaded::Bool: whether candidate optimization uses Julia threads.
• rng::AbstractRNG: random-number generator for graph sampling.
• optimizer: Optim.jl optimizer for continuous candidate parameters.
• optim_options: optional Optim.jl options object.
• optimiser: optional Optimisers.jl update rule for search parameters.
• observed: optional fixed or fitted observation model.
• alpha::Real: exponential-update coefficient used by cross-entropy search.

Returns

Return a configuration object consumed by AbstractDAGSRAlgorithm implementations.

Example

options = CommonAlgOptions(populationsize = 20, n_layers = 2)
options.populationsize == 20

init_model(alg, basis, dataset, intervals)

Construct the callable Lux model used by a differentiable symbolic-regression algorithm. basis supplies the feature count, dataset supplies target and control dimensions, and intervals contains the interval-evaluated basis values used to mask invalid inputs.

Returns

Return a model accepted by LuxCore.initialparameters, LuxCore.setup, and the call (model)(inputs, parameters, state). A custom algorithm may specialize this method when it does not use the default LayeredDAG representation.

init_cache(alg::AbstractDAGSRAlgorithm, basis, problem; kwargs...)

Build the search cache consumed by the common solve! implementation. The default method creates a Dataset, calls init_model, samples the initial population, and initializes the optimizer state. A custom algorithm may specialize this method when its cache representation differs from SearchCache.

update_parameters!(cache)

Update the population parameters for a symbolic-regression search iteration. The method is called by update_cache! after the retained candidates have been selected. Mutate the cache in place and return nothing.

convert_to_basis(candidate, parameters, options)

Convert the selected symbolic-regression candidate into a DataDrivenDiffEq.Basis. A custom graph implementation must provide this method if it does not use the package's Candidate representation.

struct AdditiveError <: DataDrivenLux.AbstractErrorModel

Additive output error model for observations following ŷ ~ y + ϵ.

When called as model(distribution, y, y_pred, scale), the model evaluates the log-likelihood of y_pred under the distribution centered at y with the supplied scale.

Returns

Return a scalar log-likelihood contribution.

struct MultiplicativeError <: DataDrivenLux.AbstractErrorModel

Multiplicative output error model for observations following ŷ ~ y * (1 + ϵ).

When called as model(distribution, y, y_pred, scale), the scale is multiplied by abs(y) before evaluating the distribution.

Returns

Return a scalar log-likelihood contribution.

struct ObservedModel{fixed, M}

The error distribution of a models output.

Construct ObservedModel(Y; fixed = false) to create one additive-normal error distribution per row of the target matrix. Set fixed = true to keep the initial scale fixed during optimization.

Arguments

• Y::AbstractMatrix: observed target data, with one target variable per row.

Keywords

• fixed::Bool: whether the observation scales are optimized.

Returns

Return an observation model used by Candidate likelihood calculations.

struct Softmax <: DataDrivenLux.AbstractSimplex

Maps an AbstractVector to the probability simplex by using softmax on each row.

Arguments

• rng::AbstractRNG: random-number generator, unused by this deterministic map.
• xhat: output buffer with the shape of x.
• x: unnormalized logits.
• kappa: positive temperature; defaults to one.

Returns

Return the output buffer after normalizing each row.

struct GumbelSoftmax <: DataDrivenLux.AbstractSimplex

Maps an AbstractVector to the probability simplex by adding gumbel distributed noise and using softmax on each row.

Arguments

• rng::AbstractRNG: random-number generator used for Gumbel noise.
• xhat: output buffer with the shape of x.
• x: unnormalized logits.
• kappa: positive temperature; defaults to one.

Returns

Return the output buffer after adding noise and normalizing each row.

Fields

struct DirectSimplex <: DataDrivenLux.AbstractSimplex

Assumes an AbstractVector is on the probability simplex.

Arguments

• rng::AbstractRNG: random-number generator, unused by this map.
• xhat: output buffer with the shape of x.
• x: probabilities that already sum to one along each row.
• kappa: accepted for interface compatibility and otherwise unused.

Returns

Return xhat after copying x into it.

Fields

struct Dataset{T, __T_x<:AbstractArray{T, 2}, __T_y<:AbstractArray{T, 2}, __T_u<:AbstractArray{T, 2}, __T_t<:AbstractArray{T, 1}, __T_x_intervals<:AbstractArray{IntervalArithmetic.Interval{T}, 1}, __T_y_intervals<:AbstractArray{IntervalArithmetic.Interval{T}, 1}, __T_u_intervals<:AbstractArray{IntervalArithmetic.Interval{T}, 1}, __T_t_interval<:IntervalArithmetic.Interval{T}}

Dataset used by DataDrivenLux candidate models, storing observed inputs, targets, controls, time points, and interval bounds for symbolic search.

Fields

• x

• y

• u

• t

• x_intervals

• y_intervals

• u_intervals

• t_interval

Arguments

• X::AbstractMatrix: state or feature data, with observations in columns.
• Y::AbstractMatrix: target data, with target variables in rows.
• U::AbstractMatrix: optional control data; defaults to an empty matrix.
• t::AbstractVector: optional time points; defaults to equally spaced indices.

Returns

Return a promoted, interval-annotated dataset used by candidate models.

struct Candidate{__T_rng, __T_st, __T_ps, __T_incoming_path, __T_outgoing_path, __T_statistics, __T_observed, __T_parameterdist, __T_scales, __T_parameters, __T_model} <: StatsAPI.StatisticalModel

A container holding all the information for the current candidate solution to the symbolic regression problem.

Fields

• rng: Random seed

• st: The current state

• ps: The current parameters

• incoming_path: Incoming paths

• outgoing_path: Outgoing path

• statistics: Statistics

• observed: The observed model

• parameterdist: The parameter distribution

• scales: The optimal scales

• parameters: The optimal parameters

• model: The component model

Arguments

• rng: random-number generator replicated for this candidate.
• model: callable graph model.
• basis: feature basis used to evaluate the dataset.
• dataset::Dataset: observed data and interval bounds.

Keywords

• observed::ObservedModel: observation likelihood model.
• parameterdist: distribution and transform for basis parameters.
• ptype: element type used for candidate state and parameters.

Returns

Return a candidate with initialized Lux parameters, path state, scales, and statistical fit values.

struct PathState{T, PO<:Tuple, PI<:Tuple} <: DataDrivenLux.AbstractPathState

State for a sampled graph path, including the current interval, applied operators, and node identifiers used to compute path complexity.

Fields

• path_interval: Accumulated loglikelihood of the state

• path_operators: All the operators of the path

• path_ids: The unique identifier of nodes in the path

Arguments

• interval::Interval: interval containing values reachable along the path.
• path_operators::Tuple: operators applied along the path.
• path_ids::Tuple: (layer, node) identifiers for the path nodes.

Returns

Return an immutable path state. Use update_path to prepend another operation without mutating an existing state.

struct FunctionNode{__T_node} <: LuxCore.AbstractLuxWrapperLayer{:node}

A layer representing a decision node with a single function and a latent array of weights representing a probability distribution over the inputs.

Fields

• node

Arguments

• f: unary or binary function represented by the node.
• arity::Int: number of inputs consumed by f.
• in_dims::Int: number of available input signals.
• id: (layer, node) identifier used in path statistics.
• input_functions: optional functions used to construct input masks.

Returns

Return a Lux wrapper layer that samples one function node configuration.

struct FunctionLayer{__T_nodes, __T_skip} <: LuxCore.AbstractLuxWrapperLayer{:nodes}

A container for multiple FunctionNodes. It accumulates all outputs of the nodes.

Fields

• nodes

• skip

Arguments

• in_dimension::Int: number of available input signals.
• arities::Tuple: arity for each function in fs.
• fs::Tuple: functions used to construct the nodes.

Keywords

• skip::Bool: include a skip connection around the layer.
• id_offset::Int: starting layer identifier for path bookkeeping.
• input_functions: optional input functions used by each node.

Returns

Return a Lux wrapper layer whose output contains the values of all nodes.

struct LayeredDAG{__T_layers} <: LuxCore.AbstractLuxWrapperLayer{:layers}

A container for a layered directed acyclic graph consisting of different FunctionLayers.

Fields

• layers

Arguments

• in_dimension::Int: number of input signals.
• out_dimension::Int: number of output equations.
• n_layers::Int: number of learned function layers.
• arities: arity for each candidate function.
• fs: candidate functions.

Keywords

• skip::Bool: retain outputs from preceding layers as inputs.
• eltype::Type: element type used for initial weights.
• input_functions: optional functions used to generate input signals.

Returns

Return a Lux wrapper model that maps candidate graph parameters and states to symbolic-regression outputs.

struct SearchCache{ALG, PTYPE, O} <: DataDrivenLux.AbstractAlgorithmCache

Optimization cache for DataDrivenLux symbolic regression algorithms.

Fields

• alg

• candidates

• ages

• keeps

• sorting

• p

• dataset

• optimiser_state

The cache owns the candidate population and the current search parameters. It is mutated by update_cache!; callers should treat it as an implementation object unless they are implementing a new AbstractDAGSRAlgorithm.

struct RelativeReward{risk} <: DataDrivenLux.AbstractRewardScale{risk}

Scales the losses in such a way that the minimum loss is equal to one.

Calling RelativeReward(risk_seeking)(losses) returns exponentially scaled rewards. With risk_seeking = true, the minimum reward is shifted to zero.

Arguments

• risk_seeking::Bool: whether to subtract the minimum reward after scaling.

struct AbsoluteReward{risk} <: DataDrivenLux.AbstractRewardScale{risk}

Scales the losses in such a way that the minimum loss is the most influential reward.

Calling AbsoluteReward(risk_seeking)(losses) uses exp.(-losses) directly. With risk_seeking = true, the minimum reward is shifted to zero.

Arguments

• risk_seeking::Bool: whether to subtract the minimum reward after scaling.

RandomSearch(
;
    populationsize,
    functions,
    arities,
    n_layers,
    skip,
    loss,
    keep,
    use_protected,
    distributed,
    threaded,
    rng,
    optimizer,
    optim_options,
    observed,
    alpha
)

Performs a random search over the space of possible solutions to the symbolic regression problem.

Keywords

The constructor accepts the fields of CommonAlgOptions: populationsize, functions, arities, n_layers, skip, loss, keep, use_protected, distributed, threaded, rng, optimizer, optim_options, observed, and alpha.

Returns

Return a AbstractDAGSRAlgorithm that updates candidate graphs by resampling without changing their continuous parameters.

Reinforce(
;
    reward,
    populationsize,
    functions,
    arities,
    n_layers,
    skip,
    loss,
    keep,
    use_protected,
    distributed,
    threaded,
    rng,
    optimizer,
    optim_options,
    observed,
    alpha,
    optimiser,
    ad_backend
)

Uses the REINFORCE algorithm to search over the space of possible solutions to the symbolic regression problem.

Keywords

• reward: RelativeReward or AbsoluteReward transform.
• ad_backend: optional DifferentiationInterface backend.
• optimiser: Optimisers.jl update rule for continuous search parameters.
• populationsize, functions, arities, n_layers, skip, loss, keep, use_protected, distributed, threaded, rng, optimizer, optim_options, observed, and alpha: forwarded to CommonAlgOptions.

Returns

Return a differentiable AbstractDAGSRAlgorithm for population search.

CrossEntropy(
;
    populationsize,
    functions,
    arities,
    n_layers,
    skip,
    loss,
    keep,
    use_protected,
    distributed,
    threaded,
    rng,
    optimizer,
    optim_options,
    observed,
    alpha
)

Uses the crossentropy method for discrete optimization to search the space of possible solutions.

Keywords

The constructor accepts populationsize, functions, arities, n_layers, skip, loss, keep, use_protected, distributed, threaded, rng, optimizer, optim_options, observed, and alpha, which are forwarded to CommonAlgOptions.

Returns

Return a AbstractDAGSRAlgorithm that updates categorical graph parameters using the cross-entropy rule.

AbstractAlgorithmCache

Developer interface for the optimization state returned by a DataDrivenLux algorithm. Concrete caches are SearchCache values and are stored in the result of solve.

AbstractSimplex

Developer interface for mappings used to normalize node weights onto the probability simplex. A subtype is called as simplex(rng, output, input, κ).

AbstractErrorModel

Developer interface for observation error models. A subtype is called as model(distribution, observation, prediction, scale) and returns a log density.

AbstractRewardScale{risk}

Developer interface for reward transformations used by search algorithms. A subtype is called with a vector of losses and returns one reward per loss.