Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: cleanup deps and replace SparseDiffTools with SciMLJacobianOperators #1124

Merged
merged 11 commits into from
Oct 7, 2024
28 changes: 12 additions & 16 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "SciMLSensitivity"
uuid = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
authors = ["Christopher Rackauckas <accounts@chrisrackauckas.com>", "Yingbo Ma <mayingbo5@gmail.com>"]
version = "7.68.0"
version = "7.69.0"

[deps]
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
Expand All @@ -13,8 +13,8 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
FunctionProperties = "f62d2435-5019-4c03-9749-2d4c77af0cbc"
Expand All @@ -24,7 +24,6 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Expand All @@ -33,19 +32,18 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e"
SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

[compat]
ADTypes = "0.1, 0.2, 1"
ADTypes = "1.9"
Accessors = "0.1.36"
Adapt = "1.0, 2.0, 3.0, 4"
AlgebraicMultigrid = "0.6.0"
Expand All @@ -56,12 +54,12 @@ ChainRulesCore = "0.10.7, 1"
ComponentArrays = "0.15.5"
DelayDiffEq = "5.43.2"
DiffEqBase = "6.151.1"
DiffEqCallbacks = "3.3"
DiffEqCallbacks = "4"
DiffEqNoiseProcess = "5.19"
Distributed = "1"
Distributions = "0.25"
EllipsisNotation = "1"
Enzyme = "0.12, 0.13"
Enzyme = "0.13"
FastBroadcast = "0.3.5"
FiniteDiff = "2"
ForwardDiff = "0.10"
FunctionProperties = "0.1"
Expand All @@ -70,15 +68,14 @@ Functors = "0.4"
GPUArraysCore = "0.1"
LinearAlgebra = "1.10"
LinearSolve = "2"
Lux = "0.5.51"
Lux = "1"
Markdown = "1.10"
ModelingToolkit = "9"
ModelingToolkit = "9.42"
NLsolve = "4.5.1"
NonlinearSolve = "3.0.1"
Optimization = "3.19.3"
OptimizationOptimisers = "0.1.6"
Optimization = "4"
OptimizationOptimisers = "0.3"
OrdinaryDiffEq = "6.81.1"
Parameters = "0.12"
Pkg = "1.10"
PreallocationTools = "0.4.4"
QuadGK = "2.9.1"
Expand All @@ -89,10 +86,10 @@ Reexport = "1.0"
ReverseDiff = "1.15.1"
SafeTestsets = "0.1.0"
SciMLBase = "2.51.4"
SciMLJacobianOperators = "0.1"
SciMLOperators = "0.3"
SciMLStructures = "1.3"
SparseArrays = "1.10"
SparseDiffTools = "2.5"
StaticArrays = "1.8.0"
StaticArraysCore = "1.4"
Statistics = "1.10"
Expand All @@ -101,7 +98,6 @@ StochasticDiffEq = "6.69"
SymbolicIndexingInterface = "0.3.26"
Test = "1.10"
Tracker = "0.2.30"
TruncatedStacktraces = "1.2"
Zygote = "0.6.67"
julia = "1.10"

Expand Down
2 changes: 1 addition & 1 deletion docs/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ Enzyme = "0.12, 0.13"
Flux = "0.14"
ForwardDiff = "0.10"
IterTools = "1"
Lux = "0.5.7, 1"
Lux = "1"
LuxCUDA = "0.3"
Optimization = "3.9, 4"
OptimizationOptimJL = "0.2, 0.3, 0.4"
Expand Down
2 changes: 1 addition & 1 deletion docs/src/examples/ode/second_order_neural.md
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ t = range(tspan[1], tspan[2], length = 20)
model = Chain(Dense(2, 50, tanh), Dense(50, 2))
ps, st = Lux.setup(Random.default_rng(), model)
ps = ComponentArray(ps)
model = Lux.Experimental.StatefulLuxLayer(model, ps, st)
model = StatefulLuxLayer{true}(model, ps, st)

ff(du, u, p, t) = model(u, p)
prob = SecondOrderODEProblem{false}(ff, du0, u0, tspan, ps)
Expand Down
89 changes: 52 additions & 37 deletions src/SciMLSensitivity.jl
Original file line number Diff line number Diff line change
@@ -1,49 +1,64 @@
module SciMLSensitivity

using DiffEqBase, ForwardDiff, Tracker, FiniteDiff, Statistics
using DiffEqCallbacks, QuadGK, RecursiveArrayTools, LinearAlgebra
import SciMLBase: AbstractNonlinearProblem
using Adapt
using LinearSolve
using Parameters: @unpack
import Accessors: @reset
import DiffEqNoiseProcess
import RandomNumbers: Xorshifts
using Random
import Zygote, ReverseDiff
import ArrayInterface
import Enzyme
import GPUArraysCore
using ADTypes
using SparseDiffTools
using SciMLOperators
using Functors
import TruncatedStacktraces
import SciMLStructures
import PreallocationTools: dualcache, get_tmp, DiffCache, LazyBufferCache,
FixedSizeDiffCache
import FunctionWrappersWrappers
using EllipsisNotation
using ADTypes: ADTypes, AutoEnzyme, AutoFiniteDiff, AutoForwardDiff,
AutoReverseDiff, AutoTracker, AutoZygote
using Accessors: @reset
using Adapt: Adapt, adapt
using ArrayInterface: ArrayInterface
using DiffEqBase: DiffEqBase, SensitivityADPassThrough
using DiffEqCallbacks: DiffEqCallbacks, IntegrandValuesSum,
IntegratingSumCallback, PresetTimeCallback
using DiffEqNoiseProcess: DiffEqNoiseProcess
using FastBroadcast: @..
using Functors: Functors, fmap
using FunctionProperties: hasbranching
using FunctionWrappersWrappers: FunctionWrappersWrappers
using GPUArraysCore: GPUArraysCore
using LinearSolve: LinearSolve
using PreallocationTools: PreallocationTools, dualcache, get_tmp, DiffCache,
FixedSizeDiffCache
using RandomNumbers: Xorshifts
using RecursiveArrayTools: RecursiveArrayTools, AbstractDiffEqArray,
AbstractVectorOfArray, ArrayPartition, DiffEqArray,
VectorOfArray
using SciMLJacobianOperators: VecJacOperator, StatefulJacobianOperator
using SciMLStructures: SciMLStructures, canonicalize, Tunable, isscimlstructure
using SymbolicIndexingInterface: SymbolicIndexingInterface, current_time, getu,
parameter_values, state_values
using QuadGK: quadgk
using SciMLBase: SciMLBase, AbstractOverloadingSensitivityAlgorithm,
AbstractForwardSensitivityAlgorithm, AbstractAdjointSensitivityAlgorithm,
AbstractSecondOrderSensitivityAlgorithm,
AbstractShadowingSensitivityAlgorithm, AbstractTimeseriesSolution,
AbstractNonlinearProblem, AbstractSensitivityAlgorithm,
AbstractDiffEqFunction, AbstractODEFunction, unwrapped_f, CallbackSet,
ContinuousCallback, DESolution, NonlinearFunction, NonlinearProblem,
DiscreteCallback, LinearProblem, ODEFunction, ODEProblem,
RODEFunction, RODEProblem, ReturnCode, SDEFunction,
SDEProblem, VectorContinuousCallback, deleteat!,
get_tmp_cache, has_adjoint, isinplace, reinit!, remake,
solve, u_modified!

using SymbolicIndexingInterface
using SciMLStructures: canonicalize, Tunable, isscimlstructure
# AD Backends
using ChainRulesCore: unthunk, @thunk, NoTangent, @not_implemented, Tangent, ZeroTangent
using Enzyme: Enzyme
using FiniteDiff: FiniteDiff
using ForwardDiff: ForwardDiff
using Tracker: Tracker, TrackedArray
using ReverseDiff: ReverseDiff
using Zygote: Zygote

using Markdown
# Std Libs
using LinearAlgebra: LinearAlgebra, Diagonal, I, UniformScaling, adjoint, axpy!,
convert, copyto!, dot, issuccess, ldiv!, lu, lu!, mul!,
norm, normalize!, qr, transpose
using Markdown: Markdown, @doc_str
using Random: Random, rand!
using Statistics: Statistics, mean

using Reexport
import ChainRulesCore: unthunk, @thunk, NoTangent, @not_implemented, Tangent, ZeroTangent
abstract type SensitivityFunction end
abstract type TransformedFunction end

import SciMLBase: unwrapped_f, _unwrap_val

import SciMLBase: AbstractOverloadingSensitivityAlgorithm, AbstractSensitivityAlgorithm,
AbstractForwardSensitivityAlgorithm, AbstractAdjointSensitivityAlgorithm,
AbstractSecondOrderSensitivityAlgorithm,
AbstractShadowingSensitivityAlgorithm,
AbstractTimeseriesSolution

include("utils.jl")
include("parameters_handling.jl")
include("sensitivity_algorithms.jl")
Expand Down
Loading
Loading