Merge branch 'master' into n-arity

Author: MilesCranmer

.github/workflows/Documentation.yml

@@ -27,8 +27,5 @@ jobs:
           DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key
         run: |
             cd docs
-            cat ../README.md | sed '/Search options/,$d' > tmp1.md
-            cat tmp1.md src/index.md > tmp2.md
-            mv tmp2.md src/index.md
             julia --project=. make.jl
 

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <[email protected]>"]
-version = "0.18.2"
+version = "0.18.6"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -15,7 +15,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
-TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 
 [weakdeps]
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
@@ -44,7 +43,6 @@ PackageExtensionCompat = "1"
 PrecompileTools = "1"
 Reexport = "1"
 SymbolicUtils = "0.19, ^1.0.5, 2"
-TestItems = "0.1"
 Zygote = "0.6"
 julia = "1.6"
 

benchmark/benchmarks.jl

@@ -1,10 +1,7 @@
 using DynamicExpressions, BenchmarkTools, Random
 
 # Trigger extensions:
-using LoopVectorization
-using Bumper
-using StrideArrays
-using Zygote
+using LoopVectorization, Bumper, StrideArrays, Zygote
 
 if PACKAGE_VERSION < v"0.14.0"
     @eval using DynamicExpressions: Node as GraphNode
@@ -18,6 +15,14 @@ else
     @eval using DynamicExpressions.NodeUtilsModule: is_constant
 end
 
+if PACKAGE_VERSION < v"0.18.6"
+    @eval using DynamicExpressions:
+        index_constants as index_constant_nodes,
+        count_constants as count_constant_nodes,
+        get_constants as get_scalar_constants,
+        set_constants! as set_scalar_constants!
+end
+
 include("../test/tree_gen_utils.jl")
 
 const SUITE = BenchmarkGroup()
@@ -113,15 +118,16 @@ end
     PACKAGE_VERSION < v"0.14.0" && return :(copy_node(t; preserve_sharing=preserve_sharing))
     return :(copy_node(t))  # Assume type used to infer sharing
 end
-@generated function get_set_constants!(tree::N) where {T,N<:AbstractExpressionNode{T}}
-    if !(@isdefined set_constants!)
-        return :(set_constants(tree, get_constants(tree)))
-    elseif hasmethod(set_constants!, Tuple{N, Vector{T}})
-        return :(set_constants!(tree, get_constants(tree)))
+@generated function get_set_constants!(tree::N) where {N}
+    T = eltype(N)
+    if !(@isdefined set_scalar_constants!)
+        return :(set_scalar_constants(tree, get_scalar_constants(tree)))
+    elseif hasmethod(set_scalar_constants!, Tuple{N, Vector{T}})
+        return :(set_scalar_constants!(tree, get_scalar_constants(tree)))
     else
         return quote
-            let (x, refs) = get_constants(tree)
-                set_constants!(tree, x, refs)
+                let (x, refs) = get_scalar_constants(tree)
+                set_scalar_constants!(tree, x, refs)
             end
         end
     end
@@ -141,12 +147,12 @@ function benchmark_utilities()
         :combine_operators,
         :count_nodes,
         :count_depth,
-        :count_constants,
+        :count_constant_nodes,
         :has_constants,
         :has_operators,
         :is_constant,
         :get_set_constants!,
-        :index_constants,
+        :index_constant_nodes,
         :string_tree,
         :hash,
     )
@@ -157,9 +163,9 @@ function benchmark_utilities()
             [
                 :simplify_tree,
                 :count_nodes,
-                :count_constants,
+                :count_constant_nodes,
                 :get_set_constants!,
-                :index_constants,
+                :index_constant_nodes,
                 :string_tree,
             ],
         )
@@ -207,7 +213,8 @@ function benchmark_utilities()
                     setup=(
                         ntrees=100;
                         n=20;
-                        trees=[$preprocess(gen_random_tree_fixed_size(n, $operators, 5, Float32)) for _ in 1:ntrees]
+                        rng=Random.MersenneTwister(0);
+                        trees=[$preprocess(gen_random_tree_fixed_size(n, $operators, 5, Float32, Node, rng)) for _ in 1:ntrees]
                     )
                 )
                 #! format: on
@@ -216,6 +223,37 @@ function benchmark_utilities()
         end
     end
 
+    # Additional methods
+    @static if PACKAGE_VERSION >= v"0.18.0"
+        suite["get_set_constants_parametric"] = @benchmarkable(
+            [get_set_constants!(ex) for ex in exs],
+            seconds = 10.0,
+            setup = (
+                operators = $operators;
+                ntrees = 100;
+                n = 20;
+                n_features = 5;
+                n_params = 3;
+                n_param_classes = 10;
+                rng = Random.MersenneTwister(0);
+                exs = [
+                    let tree = gen_random_tree_fixed_size(
+                            n, operators, n_features, Float32, ParametricNode, rng
+                        )
+                        ex = ParametricExpression(
+                            tree;
+                            operators,
+                            variable_names=map(i -> "x$i", 1:n_features),
+                            parameters=randn(rng, Float32, n_params, n_param_classes),
+                            parameter_names=map(i -> "p$i", 1:n_params),
+                        )
+                        ex
+                    end for _ in 1:ntrees
+                ]
+            )
+        )
+    end
+
     return suite
 end
 

docs/make.jl

@@ -2,6 +2,45 @@ using Documenter
 using DynamicExpressions
 using Random: AbstractRNG
 
+readme = joinpath(@__DIR__, "..", "README.md")
+
+index_content = let r = read(readme, String)
+    # Wrap img tags in raw HTML blocks:
+    r = replace(r, r"(<img\s+[^>]+>)" => s"""
+
+```@raw html
+\1
+```
+
+""")
+    # Remove end img tags:
+    r = replace(r, r"</img>" => "")
+    # Remove div tags:
+    r = replace(r, r"<div[^>]*>" => "")
+    # Remove end div tags:
+    r = replace(r, r"</div>" => "")
+
+    top_part = """
+    # Introduction
+
+    """
+
+    bottom_part = """
+    ## Contents
+
+    ```@contents
+    Pages = ["utils.md", "api.md", "eval.md"]
+    ```
+    """
+
+    join((top_part, r, bottom_part), "\n")
+end
+
+index_md = joinpath(@__DIR__, "src", "index.md")
+open(index_md, "w") do f
+    write(f, index_content)
+end
+
 makedocs(;
     sitename="DynamicExpressions.jl",
     authors="Miles Cranmer",
@@ -11,4 +50,33 @@ makedocs(;
     warnonly=true,
 )
 
+# Forward links from old docs:
+redirect_page = """
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>Redirecting...</title>
+    <script type="text/javascript">
+        var fragment = window.location.hash;
+        window.location.href = "../api/" + fragment;
+    </script>
+</head>
+<body>
+    <p>If you are not redirected automatically, follow this <a id="redirect-link" href="../api/">link to API</a>.</p>
+    <script type="text/javascript">
+        document.getElementById('redirect-link').href = "../api/" + window.location.hash;
+    </script>
+</body>
+</html>
+"""
+
+# Create the types directory and write the redirect page
+types_dir = joinpath(@__DIR__, "build", "types")
+mkpath(types_dir)
+redirect_file = joinpath(types_dir, "index.html")
+open(redirect_file, "w") do f
+    write(f, redirect_page)
+end
+
 deploydocs(; repo="github.com/SymbolicML/DynamicExpressions.jl.git")

docs/src/.gitignore

@@ -0,0 +1,2 @@
+# Gets generated:
+index.md

docs/src/types.md renamed to docs/src/api.md

@@ -1,4 +1,4 @@
-# Types
+# API
 
 ## Operator Enum
 

docs/src/index.md

@@ -1,8 +1,9 @@
 module DynamicExpressionsBumperExt
 
 using Bumper: @no_escape, @alloc
-using DynamicExpressions: OperatorEnum, AbstractExpressionNode, tree_mapreduce
-using DynamicExpressions.UtilsModule: ResultOk, counttuple, is_bad_array
+using DynamicExpressions:
+    OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array
+using DynamicExpressions.UtilsModule: ResultOk, counttuple
 
 import DynamicExpressions.ExtensionInterfaceModule:
     bumper_eval_tree_array, bumper_kern1!, bumper_kern2!
@@ -52,7 +53,7 @@ function dispatch_kerns!(operators, branch_node, cumulator, ::Val{turbo}) where
     cumulator.ok || return cumulator
 
     out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, Val(turbo))
-    return ResultOk(out, !is_bad_array(out))
+    return ResultOk(out, is_valid_array(out))
 end
 function dispatch_kerns!(
     operators, branch_node, cumulator1, cumulator2, ::Val{turbo}
@@ -63,7 +64,7 @@ function dispatch_kerns!(
     out = dispatch_kern2!(
         operators.binops, branch_node.op, cumulator1.x, cumulator2.x, Val(turbo)
     )
-    return ResultOk(out, !is_bad_array(out))
+    return ResultOk(out, is_valid_array(out))
 end
 
 @generated function dispatch_kern1!(unaops, op_idx, cumulator, ::Val{turbo}) where {turbo}

ext/DynamicExpressionsBumperExt.jl

@@ -5,8 +5,9 @@ using DynamicExpressions:
     AbstractExpressionNode,
     filter_map,
     eval_tree_array,
-    get_constants,
-    set_constants!
+    get_scalar_constants,
+    set_scalar_constants!,
+    get_number_type
 using Compat: @inline
 
 import Optim: Optim, OptimizationResults, NLSolversBase
@@ -44,9 +45,14 @@ function wrap_func(
     function wrapped_f(args::Vararg{Any,M}) where {M}
         first_args = args[begin:(end - 1)]
         x = args[end]
-        set_constants!(tree, x, refs)
+        set_scalar_constants!(tree, x, refs)
         return @inline(f(first_args..., tree))
     end
+    # without first args, it looks like this
+    # function wrapped_f(x)
+    #     set_scalar_constants!(tree, x, refs)
+    #     return @inline(f(tree))
+    # end
     return wrapped_f
 end
 function wrap_func(
@@ -100,7 +106,8 @@ function Optim.optimize(
     if make_copy
         tree = copy(tree)
     end
-    x0, refs = get_constants(tree)
+
+    x0, refs = get_scalar_constants(tree)
     if !isnothing(h!)
         throw(
             ArgumentError(
@@ -117,7 +124,7 @@ function Optim.optimize(
         )
     end
     minimizer = Optim.minimizer(base_res)
-    set_constants!(tree, minimizer, refs)
+    set_scalar_constants!(tree, minimizer, refs)
     return ExpressionOptimizationResults(base_res, tree)
 end