From f1cf5329b0457c5a116d0d8eba8b0a399ae82e7f Mon Sep 17 00:00:00 2001 From: Olivier Cots Date: Tue, 19 May 2026 22:11:32 +0200 Subject: [PATCH] Add HamiltonianSystem with AD support and AugmentedHamiltonianTrait - Add AugmentedHamiltonianTrait and AbstractAugmentedHamiltonianConfig in Common module - Update AbstractHamiltonianSystem to include AD trait parameter (AT) - Add HamiltonianSystem struct with AD backend support, mirroring HamiltonianVectorFieldSystem - Implement build_rhs_augmented with batch compatibility checks for matrix inputs - Add build_system overloads for Data.Hamiltonian with AD backend - Update all references to AbstractHamiltonianSystem to include AT parameter: - AbstractHamiltonianFlow in Flows - MultiPhaseHamiltonianFlow in MultiPhase - All test files with fake Hamiltonian systems - Add comprehensive docstrings following documentation standards - Add tests for HamiltonianSystem, ad_trait, build_rhs_augmented, and type stability - All 1786 tests passing --- src/Common/Common.jl | 4 +- src/Common/abstract_trait.jl | 14 + src/Common/configs.jl | 19 ++ src/Flows/abstract_flow.jl | 2 +- src/MultiPhase/multiphase_flow.jl | 6 +- src/Systems/Systems.jl | 5 + src/Systems/abstract_system.jl | 25 +- src/Systems/building.jl | 51 ++++ src/Systems/hamiltonian_system.jl | 196 ++++++++++++++ .../hamiltonian_vector_field_system.jl | 2 +- test/suite/common/test_abstract_trait.jl | 30 ++- test/suite/flows/test_calling_flows.jl | 2 +- test/suite/flows/test_flow_callables.jl | 2 +- .../multiphase/test_calling_multiphase.jl | 2 +- test/suite/multiphase/test_concatenation.jl | 2 +- test/suite/systems/test_abstract_system.jl | 2 +- test/suite/systems/test_hamiltonian_system.jl | 252 ++++++++++++++++++ .../test_hamiltonian_vector_field_system.jl | 13 +- 18 files changed, 614 insertions(+), 15 deletions(-) create mode 100644 src/Systems/hamiltonian_system.jl create mode 100644 test/suite/systems/test_hamiltonian_system.jl diff --git a/src/Common/Common.jl b/src/Common/Common.jl index 2a0cebe2..875c8d4e 100644 --- a/src/Common/Common.jl +++ b/src/Common/Common.jl @@ -40,11 +40,11 @@ include(joinpath(@__DIR__, "internal_norm.jl")) export AbstractTag export AbstractTrait, AbstractModeTrait, AbstractContentTrait, AbstractMutabilityTrait, AbstractADTrait -export PointTrait, TrajectoryTrait, StateTrait, HamiltonianTrait +export PointTrait, TrajectoryTrait, StateTrait, HamiltonianTrait, AugmentedHamiltonianTrait export InPlace, OutOfPlace export WithAD, WithoutAD export AbstractCache -export AbstractConfig, AbstractPointConfig, AbstractTrajectoryConfig, AbstractStateConfig, AbstractHamiltonianConfig +export AbstractConfig, AbstractPointConfig, AbstractTrajectoryConfig, AbstractStateConfig, AbstractHamiltonianConfig, AbstractAugmentedHamiltonianConfig export StatePointConfig, StateTrajectoryConfig, HamiltonianPointConfig, HamiltonianTrajectoryConfig export tspan, initial_condition, initial_state, initial_costate export VariableDependence, Fixed, NonFixed diff --git a/src/Common/abstract_trait.jl b/src/Common/abstract_trait.jl index 59fed0c4..9890a2a0 100644 --- a/src/Common/abstract_trait.jl +++ b/src/Common/abstract_trait.jl @@ -335,6 +335,20 @@ See also: [`CTFlows.Common.StateTrait`](@ref), [`CTFlows.Common.AbstractContentT """ struct HamiltonianTrait <: AbstractContentTrait end +""" +$(TYPEDEF) + +Trait marker for augmented Hamiltonian systems, where the Hamiltonian includes an augmented variable (e.g., a parameter or control variable) in addition to state and costate variables. + +# Notes +- Used in conjunction with [`CTFlows.Common.AbstractAugmentedHamiltonianConfig`](@ref) to specify that a configuration is for an augmented Hamiltonian system. +- Subtypes [`CTFlows.Common.AbstractContentTrait`](@ref). +- Used to distinguish augmented Hamiltonian systems from standard Hamiltonian systems in trait-based dispatch. + +See also: [`CTFlows.Common.HamiltonianTrait`](@ref), [`CTFlows.Common.AbstractAugmentedHamiltonianConfig`](@ref), [`CTFlows.Common.AbstractContentTrait`](@ref). +""" +struct AugmentedHamiltonianTrait <: AbstractContentTrait end + # ============================================================================= # AD trait family (automatic differentiation capability) # ============================================================================= diff --git a/src/Common/configs.jl b/src/Common/configs.jl index c73468f3..7434d79f 100644 --- a/src/Common/configs.jl +++ b/src/Common/configs.jl @@ -190,6 +190,25 @@ Matches any `AbstractConfig` with `HamiltonianTrait` as the content parameter. """ const AbstractHamiltonianConfig{X0, M} = AbstractConfigWithMaC{X0, M, HamiltonianTrait} +""" +$(TYPEDEF) + +Type alias for augmented Hamiltonian configurations, which include state, costate, and augmented variable initial conditions. + +# Type Parameters +- `X0`: Type of the initial condition (typically a vector concatenating state, costate, and augmented variable). +- `M`: Type of the mutability trait (`InPlace` or `OutOfPlace`). + +# Notes +- Augmented Hamiltonian configurations are used for systems where the Hamiltonian depends on an additional variable (e.g., a control parameter or optimization variable). +- The initial condition typically has the form `vcat(x0, p0, pv0)` where `x0` is the initial state, `p0` is the initial costate, and `pv0` is the initial augmented variable. +- Subtypes [`CTFlows.Common.AbstractConfigWithMaC`](@ref) with [`CTFlows.Common.AugmentedHamiltonianTrait`](@ref). +- Used in conjunction with [`CTFlows.Systems.HamiltonianSystem`](@ref) for automatic differentiation-based Hamiltonian integration. + +See also: [`CTFlows.Common.AbstractHamiltonianConfig`](@ref), [`CTFlows.Common.AugmentedHamiltonianTrait`](@ref), [`CTFlows.Systems.HamiltonianSystem`](@ref). +""" +const AbstractAugmentedHamiltonianConfig{X0, M} = AbstractConfigWithMaC{X0, M, AugmentedHamiltonianTrait} + # ============================================================================= # Interface implementations on abstract config types # ============================================================================= diff --git a/src/Flows/abstract_flow.jl b/src/Flows/abstract_flow.jl index 5e41eb1a..89b58fe7 100644 --- a/src/Flows/abstract_flow.jl +++ b/src/Flows/abstract_flow.jl @@ -82,7 +82,7 @@ true See also: [`CTFlows.Flows.AbstractFlow`](@ref), [`CTFlows.Flows.AbstractStateFlow`](@ref), [`CTFlows.Systems.AbstractHamiltonianSystem`](@ref). """ -abstract type AbstractHamiltonianFlow{TD, VD, S<:Systems.AbstractHamiltonianSystem{TD,VD}} <: AbstractFlow{TD, VD} end +abstract type AbstractHamiltonianFlow{TD, VD, S<:Systems.AbstractHamiltonianSystem{TD,VD,<:Common.AbstractADTrait}} <: AbstractFlow{TD, VD} end """ $(TYPEDSIGNATURES) diff --git a/src/MultiPhase/multiphase_flow.jl b/src/MultiPhase/multiphase_flow.jl index ac812215..910512ea 100644 --- a/src/MultiPhase/multiphase_flow.jl +++ b/src/MultiPhase/multiphase_flow.jl @@ -71,9 +71,9 @@ mpf = MultiPhaseHamiltonianFlow([flow1, flow2], [1.0], [nothing]) See also: [`CTFlows.MultiPhase.MultiPhaseStateFlow`](@ref), [`CTFlows.Flows.HamiltonianFlow`](@ref). """ struct MultiPhaseHamiltonianFlow{ - TD<:Common.TimeDependence, - VD<:Common.VariableDependence, - S<:Systems.AbstractHamiltonianSystem{TD, VD}, + TD<:Common.TimeDependence, + VD<:Common.VariableDependence, + S<:Systems.AbstractHamiltonianSystem{TD, VD, <:Common.AbstractADTrait}, I<:Integrators.AbstractIntegrator, ST<:Vector{<:Real}, J<:Vector{<:Any}} <: Flows.AbstractHamiltonianFlow{TD, VD, S} diff --git a/src/Systems/Systems.jl b/src/Systems/Systems.jl index db392886..6c522377 100644 --- a/src/Systems/Systems.jl +++ b/src/Systems/Systems.jl @@ -19,6 +19,7 @@ import CTBase.Exceptions using ..Common using ..Data +using ..Differentiation # ============================================================================== # Include files @@ -27,6 +28,7 @@ using ..Data include(joinpath(@__DIR__, "abstract_system.jl")) include(joinpath(@__DIR__, "vector_field_system.jl")) include(joinpath(@__DIR__, "hamiltonian_vector_field_system.jl")) +include(joinpath(@__DIR__, "hamiltonian_system.jl")) include(joinpath(@__DIR__, "building.jl")) # ============================================================================== @@ -39,6 +41,9 @@ export rhs_oop export state_dimension export VectorFieldSystem export HamiltonianVectorFieldSystem +export HamiltonianSystem export build_system +export build_rhs_augmented +export ad_trait end # module Systems diff --git a/src/Systems/abstract_system.jl b/src/Systems/abstract_system.jl index cf9f0cec..3272e497 100644 --- a/src/Systems/abstract_system.jl +++ b/src/Systems/abstract_system.jl @@ -68,6 +68,7 @@ Carries the time-dependence and variable-dependence traits for compile-time disp # Type Parameters - `TD <: TimeDependence`: Time dependence trait (Autonomous or NonAutonomous) - `VD <: VariableDependence`: Variable dependence trait (Fixed or NonFixed) +- `AT <: AbstractADTrait`: AD capability trait (WithAD or WithoutAD) # Example \`\`\`julia-repl @@ -79,7 +80,7 @@ true See also: [`CTFlows.Systems.AbstractSystem`](@ref), [`CTFlows.Systems.AbstractStateSystem`](@ref). """ -abstract type AbstractHamiltonianSystem{TD, VD} <: AbstractSystem{TD, VD} end +abstract type AbstractHamiltonianSystem{TD, VD, AT<:Common.AbstractADTrait} <: AbstractSystem{TD, VD} end """ $(TYPEDSIGNATURES) @@ -194,6 +195,28 @@ end """ $(TYPEDSIGNATURES) +Return the automatic differentiation capability trait of a Hamiltonian system. + +# Arguments +- `sys::AbstractHamiltonianSystem`: The Hamiltonian system. + +# Returns +- `AT <: AbstractADTrait`: The AD capability trait, either [`CTFlows.Common.WithAD`](@ref) or [`CTFlows.Common.WithoutAD`](@ref). + +# Notes +- [`CTFlows.Systems.HamiltonianVectorFieldSystem`](@ref) always returns `WithoutAD` since it uses an explicitly provided vector field. +- [`CTFlows.Systems.HamiltonianSystem`](@ref) returns `WithAD` since it uses automatic differentiation to compute gradients from a scalar Hamiltonian function. +- This trait is used for dispatch in flow integration and cache preparation. + +See also: [`CTFlows.Common.AbstractADTrait`](@ref), [`CTFlows.Common.WithAD`](@ref), [`CTFlows.Common.WithoutAD`](@ref), [`CTFlows.Systems.HamiltonianVectorFieldSystem`](@ref), [`CTFlows.Systems.HamiltonianSystem`](@ref). +""" +function ad_trait(sys::AbstractHamiltonianSystem{TD, VD, AT}) where {TD, VD, AT} + return AT +end + +""" +$(TYPEDSIGNATURES) + Return the right-hand side function for the system. The returned function must have the signature `(du, u, p, t) -> nothing` and diff --git a/src/Systems/building.jl b/src/Systems/building.jl index cd6deaa3..1985131f 100644 --- a/src/Systems/building.jl +++ b/src/Systems/building.jl @@ -83,3 +83,54 @@ function build_system(hvf::Data.HamiltonianVectorField; state_dimension::Union{I return HamiltonianVectorFieldSystem(hvf; state_dimension=state_dimension) end +""" +$(TYPEDSIGNATURES) + +Build a [`CTFlows.Systems.HamiltonianSystem`](@ref) from a scalar `Hamiltonian` function with automatic differentiation. + +Constructs a concrete Hamiltonian system that wraps the scalar Hamiltonian function with an AD backend to compute gradients on-the-fly. The resulting system is ready for use with flow integration pipelines. + +# Arguments +- `h::Data.Hamiltonian`: The scalar Hamiltonian function to wrap into a system. +- `backend::Differentiation.AbstractADBackend`: The automatic differentiation backend (e.g., `AutoForwardDiff`, `AutoZygote`). +- `state_dimension::Union{Int, Nothing}`: The state dimension (number of state variables, not including costates). Defaults to `nothing` (inferred at runtime). + +# Returns +- `HamiltonianSystem`: A concrete Hamiltonian system with automatic differentiation support. + +# Example +\`\`\`julia-repl +julia> using CTFlows.Systems, CTFlows.Common, CTFlows.Data + +julia> h = Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2); autonomous=true, variable=false) +Hamiltonian{var"#1", Autonomous, Fixed} + +julia> sys = build_system(h, AutoForwardDiff()) +HamiltonianSystem + time_dependence: Autonomous + variable_dependence: Fixed + state_dimension: unknown + hamiltonian: Hamiltonian{var"#1", Autonomous, Fixed} + backend: AutoForwardDiff() + +julia> sys = build_system(h, AutoForwardDiff(); state_dimension=3) +HamiltonianSystem + time_dependence: Autonomous + variable_dependence: Fixed + state_dimension: 3 + hamiltonian: Hamiltonian{var"#1", Autonomous, Fixed} + backend: AutoForwardDiff() +\`\`\` + +# Notes +- The AD backend is used to compute Hamiltonian gradients `∂H/∂x` and `∂H/∂p` automatically during integration. +- Specifying `state_dimension` improves type stability and performance by closing dimensions in the pre-computed RHS closures. +- This overload is for scalar Hamiltonian functions where gradients are computed via AD. For explicit vector fields, use [`CTFlows.Systems.HamiltonianVectorFieldSystem`](@ref) instead. + +See also: [`CTFlows.Data.Hamiltonian`](@ref), [`CTFlows.Systems.HamiltonianSystem`](@ref), [`CTFlows.Systems.HamiltonianVectorFieldSystem`](@ref), [`CTFlows.Differentiation.AbstractADBackend`](@ref). +""" +function build_system(h::Data.Hamiltonian, backend::Differentiation.AbstractADBackend; + state_dimension::Union{Int,Nothing}=Common.__state_dimension()) + return HamiltonianSystem(h, backend; state_dimension=state_dimension) +end + diff --git a/src/Systems/hamiltonian_system.jl b/src/Systems/hamiltonian_system.jl new file mode 100644 index 00000000..d443e679 --- /dev/null +++ b/src/Systems/hamiltonian_system.jl @@ -0,0 +1,196 @@ +# ============================================================================= +# HamiltonianSystem — AD-based Hamiltonian system with scalar Hamiltonian function +# ============================================================================= + +""" +$(TYPEDEF) + +Concrete `AbstractHamiltonianSystem` wrapping a scalar `Hamiltonian` function with an AD backend. +The state dimension `N` is stored as a type parameter for compile-time validation and performance. +If `N=nothing`, the dimension is inferred at runtime. + +# Type Parameters +- `N`: State dimension (`Int` if known, `nothing` if unknown). +- `F`: concrete type of the wrapped Hamiltonian function. +- `TD <: TimeDependence`: `Autonomous` or `NonAutonomous`. +- `VD <: VariableDependence`: `Fixed` or `NonFixed`. +- `BACKEND <: AbstractADBackend`: concrete AD backend type. +- `RHS`: type of the pre-computed in-place right-hand side function. +- `OOPROHS`: type of the pre-computed out-of-place right-hand side function. + +# Fields +- `h::Hamiltonian{F, TD, VD}`: the underlying scalar Hamiltonian function. +- `backend::BACKEND`: the AD backend for gradient computation. +- `rhs::RHS`: the pre-computed in-place right-hand side closure with signature `(du, u, p, t) -> nothing`. +- `rhs_oop::OOPROHS`: the pre-computed out-of-place right-hand side closure with signature `(u, p, t) -> du`. + +# Example +```julia-repl +julia> using CTFlows.Systems, CTFlows.Common, CTFlows.Data + +julia> h = Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2); autonomous=true, variable=false) +Hamiltonian{var"#1", Autonomous, Fixed} + +julia> sys = HamiltonianSystem(h, AutoForwardDiff(); state_dimension=3) +HamiltonianSystem + time_dependence: Autonomous + variable_dependence: Fixed + state_dimension: 3 + hamiltonian: Hamiltonian{var"#1", Autonomous, Fixed} + backend: AutoForwardDiff() +``` + +See also: [`CTFlows.Data.Hamiltonian`](@ref), [`CTFlows.Systems.AbstractHamiltonianSystem`](@ref), [`CTFlows.Common.AbstractADTrait`](@ref). +""" +struct HamiltonianSystem{ + N, + F<:Function, + TD<:Common.TimeDependence, + VD<:Common.VariableDependence, + BACKEND<:Differentiation.AbstractADBackend, + RHS<:Function, + OOPROHS<:Function, +} <: AbstractHamiltonianSystem{TD, VD, Common.WithAD} + h::Data.Hamiltonian{F, TD, VD} + backend::BACKEND + rhs::RHS + rhs_oop::OOPROHS +end + +# ============================================================================= +# Constructors +# ============================================================================= + +function HamiltonianSystem(h::Data.Hamiltonian{F,TD,VD}, backend::Differentiation.AbstractADBackend; + state_dimension::Union{Int,Nothing}=Common.__state_dimension()) where {F,TD,VD} + rhs = _build_rhs_hs(h, backend, Val(state_dimension)) + rhs_oop = _build_oop_rhs_hs(h, backend, Val(state_dimension)) + return HamiltonianSystem{state_dimension, F, TD, VD, typeof(backend), + typeof(rhs), typeof(rhs_oop)}(h, backend, rhs, rhs_oop) +end + +# ============================================================================= +# Internal helpers for augmented split/assign (Vector + Matrix, concrete integers) +# ============================================================================= + +_aug_split(u::AbstractVector, n_x::Int, n_v::Int) = + (@view(u[1:n_x]), @view(u[n_x+1:2*n_x]), @view(u[end-n_v+1:end])) +_aug_split(u::AbstractMatrix, n_x::Int, n_v::Int) = + (@view(u[1:n_x,:]), @view(u[n_x+1:2*n_x,:]), @view(u[end-n_v+1:end,:])) + +_aug_assign!(du::AbstractVector, dx, dp, dpv, n_x::Int, n_v::Int) = + (du[1:n_x] .= dx; du[n_x+1:2*n_x] .= dp; du[end-n_v+1:end] .= dpv) +_aug_assign!(du::AbstractMatrix, dx, dp, dpv, n_x::Int, n_v::Int) = + (du[1:n_x,:] .= dx; du[n_x+1:2*n_x,:] .= dp; du[end-n_v+1:end,:] .= dpv) + +# ============================================================================= +# Internal helpers for building RHS (in-place) +# ============================================================================= + +function _build_rhs_hs(h, backend, ::Val{N}) where N + return function (du, u, λ, t) + x, p = _ham_split(u, N) + ∂x, ∂p = Differentiation.hamiltonian_gradient(backend, h, t, x, p, Common.variable(λ), Common.cache(λ)) + _ham_assign!(du, ∂p, -∂x, N) + return nothing + end +end + +# ============================================================================= +# Internal helpers for building RHS (out-of-place) +# ============================================================================= + +function _build_oop_rhs_hs(h, backend, ::Val{N}) where N + return function (u, λ, t) + x, p = _ham_split(u, N) + ∂x, ∂p = Differentiation.hamiltonian_gradient(backend, h, t, x, p, Common.variable(λ), Common.cache(λ)) + return vcat(∂p, -∂x) + end +end + +# ============================================================================= +# Batch compatibility check for augmented RHS +# ============================================================================= + +function _check_aug_batch_compat(u::AbstractMatrix, v::AbstractMatrix) + if size(u, 2) != size(v, 2) + throw(Exceptions.PreconditionError( + "batch size mismatch in augmented Hamiltonian RHS"; + reason = "size(u, 2) = $(size(u, 2)) ≠ size(v, 2) = $(size(v, 2))", + context = "build_rhs_augmented — matrix batch mode", + suggestion = "variable v must have the same number of columns as the state u", + )) + end + return nothing +end +_check_aug_batch_compat(u, v) = nothing # no-op for non-matrix cases + +# ============================================================================= +# build_rhs_augmented — lazy, closes over concrete n_x and n_v +# ============================================================================= + +function build_rhs_augmented(sys::HamiltonianSystem, n_x::Int, n_v::Int) + h, backend = sys.h, sys.backend + return function (du, u, λ, t) + v = Common.variable(λ) + _check_aug_batch_compat(u, v) # no-op if not matrix or matrix compatible + x, p, pv = _aug_split(u, n_x, n_v) + ∂x, ∂p = Differentiation.hamiltonian_gradient(backend, h, t, x, p, v, Common.cache(λ)) + ∂v = Differentiation.variable_gradient(backend, h, t, x, p, v, Common.cache(λ)) + _aug_assign!(du, ∂p, -∂x, -∂v, n_x, n_v) + return nothing + end +end + +# ============================================================================= +# rhs accessor (in-place) +# ============================================================================= + +function rhs(sys::HamiltonianSystem) + return sys.rhs +end + +# ============================================================================= +# rhs_oop accessor (out-of-place) +# ============================================================================= + +function rhs_oop(sys::HamiltonianSystem, is_u0_mutable::Bool=true) + return sys.rhs_oop +end + +# ============================================================================= +# state_dimension accessor +# ============================================================================= + +function state_dimension(sys::HamiltonianSystem{N}) where N + return N +end + +# ============================================================================= +# Validation +# ============================================================================= + +function _check_state_dimension(sys::HamiltonianSystem{N}, x0) where N + if N !== nothing && length(x0) != N + throw(Exceptions.PreconditionError( + "state dimension mismatch"; + reason = "length(x0) = $(length(x0)) ≠ N = $N", + context = "HamiltonianSystem construction with known dimension", + suggestion = "either omit state_dimension or ensure length(x0) matches the specified dimension", + )) + end + return nothing +end + +# ============================================================================= +# Base.show +# ============================================================================= + +function Base.show(io::IO, sys::HamiltonianSystem) + println(io, "HamiltonianSystem") + println(io, " time_dependence: ", Common.time_dependence(sys)) + println(io, " variable_dependence: ", Common.variable_dependence(sys)) + println(io, " state_dimension: ", sys.N === nothing ? "unknown" : sys.N) + println(io, " hamiltonian: ", sys.h) + println(io, " backend: ", sys.backend) +end diff --git a/src/Systems/hamiltonian_vector_field_system.jl b/src/Systems/hamiltonian_vector_field_system.jl index 2e669c26..4e6c5abc 100644 --- a/src/Systems/hamiltonian_vector_field_system.jl +++ b/src/Systems/hamiltonian_vector_field_system.jl @@ -51,7 +51,7 @@ HamiltonianVectorFieldSystem See also: [`CTFlows.Data.HamiltonianVectorField`](@ref), [`CTFlows.Systems.AbstractHamiltonianSystem`](@ref), [`CTFlows.Common.TimeDependence`](@ref), [`CTFlows.Common.VariableDependence`](@ref). """ -struct HamiltonianVectorFieldSystem{N, F<:Function, TD<:Common.TimeDependence, VD<:Common.VariableDependence, MD<:Common.AbstractMutabilityTrait, RHS<:Function, OOPROHS<:Function, FINRHS} <: AbstractHamiltonianSystem{TD, VD} +struct HamiltonianVectorFieldSystem{N, F<:Function, TD<:Common.TimeDependence, VD<:Common.VariableDependence, MD<:Common.AbstractMutabilityTrait, RHS<:Function, OOPROHS<:Function, FINRHS} <: AbstractHamiltonianSystem{TD, VD, Common.WithoutAD} hvf::Data.HamiltonianVectorField{F, TD, VD, MD} rhs::RHS rhs_oop::OOPROHS diff --git a/test/suite/common/test_abstract_trait.jl b/test/suite/common/test_abstract_trait.jl index e9d79e10..8bc432a1 100644 --- a/test/suite/common/test_abstract_trait.jl +++ b/test/suite/common/test_abstract_trait.jl @@ -176,6 +176,25 @@ function test_abstract_trait() end end + Test.@testset "AugmentedHamiltonianTrait" begin + Test.@testset "AugmentedHamiltonianTrait is exported" begin + Test.@test isdefined(Common, :AugmentedHamiltonianTrait) + end + + Test.@testset "AugmentedHamiltonianTrait is concrete" begin + Test.@test !isabstracttype(Common.AugmentedHamiltonianTrait) + end + + Test.@testset "AugmentedHamiltonianTrait instantiates" begin + aug = Common.AugmentedHamiltonianTrait() + Test.@test aug isa Common.AugmentedHamiltonianTrait + end + + Test.@testset "AugmentedHamiltonianTrait subtypes AbstractContentTrait" begin + Test.@test Common.AugmentedHamiltonianTrait <: Common.AbstractContentTrait + end + end + Test.@testset "InPlace" begin Test.@testset "InPlace is exported" begin Test.@test isdefined(Common, :InPlace) @@ -263,6 +282,7 @@ function test_abstract_trait() Test.@test Common.TrajectoryTrait <: Common.AbstractTrait Test.@test Common.StateTrait <: Common.AbstractTrait Test.@test Common.HamiltonianTrait <: Common.AbstractTrait + Test.@test Common.AugmentedHamiltonianTrait <: Common.AbstractTrait Test.@test Common.InPlace <: Common.AbstractTrait Test.@test Common.OutOfPlace <: Common.AbstractTrait Test.@test Common.WithAD <: Common.AbstractTrait @@ -274,6 +294,7 @@ function test_abstract_trait() Test.@test !(Common.TrajectoryTrait <: Common.AbstractContentTrait) Test.@test !(Common.StateTrait <: Common.AbstractModeTrait) Test.@test !(Common.HamiltonianTrait <: Common.AbstractModeTrait) + Test.@test !(Common.AugmentedHamiltonianTrait <: Common.AbstractModeTrait) end end @@ -284,11 +305,18 @@ function test_abstract_trait() Test.@testset "Exports Verification" begin Test.@testset "Exported trait types" begin for sym in (:AbstractTrait, :AbstractModeTrait, :AbstractContentTrait, :AbstractMutabilityTrait, :AbstractADTrait, - :PointTrait, :TrajectoryTrait, :StateTrait, :HamiltonianTrait, + :PointTrait, :TrajectoryTrait, :StateTrait, :HamiltonianTrait, :AugmentedHamiltonianTrait, :InPlace, :OutOfPlace, :WithAD, :WithoutAD, :AbstractCache) Test.@test isdefined(Common, sym) end end + + Test.@testset "Exported config aliases" begin + for sym in (:AbstractConfig, :AbstractPointConfig, :AbstractTrajectoryConfig, :AbstractStateConfig, + :AbstractHamiltonianConfig, :AbstractAugmentedHamiltonianConfig) + Test.@test isdefined(Common, sym) + end + end end end end diff --git a/test/suite/flows/test_calling_flows.jl b/test/suite/flows/test_calling_flows.jl index 77fbf8ba..ed75b009 100644 --- a/test/suite/flows/test_calling_flows.jl +++ b/test/suite/flows/test_calling_flows.jl @@ -24,7 +24,7 @@ end """ Fake Hamiltonian system for testing the calling workflow. """ -struct FakeHamiltonianSystemForCalling <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed} +struct FakeHamiltonianSystemForCalling <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed, Common.WithoutAD} state_dim::Int end diff --git a/test/suite/flows/test_flow_callables.jl b/test/suite/flows/test_flow_callables.jl index 5a09c9b8..6a5a27c0 100644 --- a/test/suite/flows/test_flow_callables.jl +++ b/test/suite/flows/test_flow_callables.jl @@ -20,7 +20,7 @@ struct FakeStateSystemFC <: Systems.AbstractStateSystem{Common.Autonomous, Commo n::Int end -struct FakeHamSysFC <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed} +struct FakeHamSysFC <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed, Common.WithoutAD} n::Int end diff --git a/test/suite/multiphase/test_calling_multiphase.jl b/test/suite/multiphase/test_calling_multiphase.jl index b869d4c4..19c8f6c7 100644 --- a/test/suite/multiphase/test_calling_multiphase.jl +++ b/test/suite/multiphase/test_calling_multiphase.jl @@ -23,7 +23,7 @@ function Systems.rhs(sys::FakeStateSystem) return (du, u, p, t) -> du .= sys.data .* u end -struct FakeHamiltonianSystem <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed} +struct FakeHamiltonianSystem <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed, Common.WithoutAD} data::Vector{Float64} end diff --git a/test/suite/multiphase/test_concatenation.jl b/test/suite/multiphase/test_concatenation.jl index e194c023..7d8ed368 100644 --- a/test/suite/multiphase/test_concatenation.jl +++ b/test/suite/multiphase/test_concatenation.jl @@ -23,7 +23,7 @@ end Systems.rhs(::FakeStateSystem) = (du, u, p, t) -> (du .= u) -struct FakeHamiltonianSystem <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed} +struct FakeHamiltonianSystem <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed, Common.WithoutAD} tag::Symbol end diff --git a/test/suite/systems/test_abstract_system.jl b/test/suite/systems/test_abstract_system.jl index 16a5ef9b..fc5d6e17 100644 --- a/test/suite/systems/test_abstract_system.jl +++ b/test/suite/systems/test_abstract_system.jl @@ -36,7 +36,7 @@ function Systems.rhs(sys::FakeStateSystem) return (du, u, p, t) -> du .= sys.data .* u end -struct FakeHamiltonianSystem <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed} +struct FakeHamiltonianSystem <: Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed, Common.WithoutAD} data::Vector{Float64} end diff --git a/test/suite/systems/test_hamiltonian_system.jl b/test/suite/systems/test_hamiltonian_system.jl new file mode 100644 index 00000000..c1604301 --- /dev/null +++ b/test/suite/systems/test_hamiltonian_system.jl @@ -0,0 +1,252 @@ +module TestHamiltonianSystem + +import Test +import CTBase.Exceptions +import CTFlows.Data: Data +import CTFlows.Common: Common +import CTFlows.Systems: Systems +import CTFlows.Differentiation: Differentiation + +const VERBOSE = isdefined(Main, :TestData) ? Main.TestData.VERBOSE : true +const SHOWTIMING = isdefined(Main, :TestData) ? Main.TestData.SHOWTIMING : true + +# Fake AD backend for testing (no actual AD, just a stub) +struct FakeADBackend <: Differentiation.AbstractADBackend end + +function Differentiation.hamiltonian_gradient(backend::FakeADBackend, h, t, x, p, v, cache) + # Fake gradient: ∂H/∂x = x, ∂H/∂p = p (for H = 0.5*(x^2 + p^2)) + return (x, p) +end + +function Differentiation.variable_gradient(backend::FakeADBackend, h, t, x, p, v, cache) + # Fake gradient: ∂H/∂v = v (for H = 0.5*v^2), or 0.0 if v is nothing + return v === nothing ? 0.0 : v +end + +function test_hamiltonian_system() + Test.@testset "Hamiltonian System Tests" verbose=VERBOSE showtiming=SHOWTIMING begin + + # ==================================================================== + # UNIT TESTS - Construction + # ==================================================================== + + Test.@testset "Construction" begin + h = Data.Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2); is_autonomous=true, is_variable=false) + backend = FakeADBackend() + + # Without state_dimension + sys1 = Systems.HamiltonianSystem(h, backend) + Test.@test sys1 isa Systems.HamiltonianSystem + Test.@test sys1 isa Systems.AbstractHamiltonianSystem + Test.@test Systems.state_dimension(sys1) === nothing + + # With state_dimension + sys2 = Systems.HamiltonianSystem(h, backend; state_dimension=2) + Test.@test sys2 isa Systems.HamiltonianSystem + Test.@test Systems.state_dimension(sys2) == 2 + end + + # ==================================================================== + # UNIT TESTS - ad_trait + # ==================================================================== + + Test.@testset "ad_trait" begin + h = Data.Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2); is_autonomous=true, is_variable=false) + backend = FakeADBackend() + sys = Systems.HamiltonianSystem(h, backend) + + Test.@test Systems.ad_trait(sys) === Common.WithAD + end + + # ==================================================================== + # UNIT TESTS - rhs (in-place) + # ==================================================================== + + Test.@testset "rhs" begin + h = Data.Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2); is_autonomous=true, is_variable=false) + backend = FakeADBackend() + sys = Systems.HamiltonianSystem(h, backend; state_dimension=2) + + rhs = Systems.rhs(sys) + Test.@test rhs isa Function + + # Test RHS call with vector + u = [1.0, 2.0, 3.0, 4.0] # x = [1, 2], p = [3, 4] + du = zeros(4) + p = Common.ODEParameters(nothing) + rhs(du, u, p, 0.0) + + # ∂H/∂x = x = [1, 2], ∂H/∂p = p = [3, 4], so du = [∂p, -∂x] = [3, 4, -1, -2] + Test.@test du[1:2] == [3.0, 4.0] + Test.@test du[3:4] == [-1.0, -2.0] + + # Test RHS call with matrix + u_mat = [1.0 2.0; 3.0 4.0; 5.0 6.0; 7.0 8.0] # x = [1 2; 3 4], p = [5 6; 7 8] + du_mat = zeros(4, 2) + rhs(du_mat, u_mat, p, 0.0) + Test.@test du_mat ≈ [5.0 6.0; 7.0 8.0; -1.0 -2.0; -3.0 -4.0] atol=1e-10 + end + + # ==================================================================== + # UNIT TESTS - rhs_oop (out-of-place) + # ==================================================================== + + Test.@testset "rhs_oop" begin + h = Data.Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2); is_autonomous=true, is_variable=false) + backend = FakeADBackend() + sys = Systems.HamiltonianSystem(h, backend; state_dimension=2) + + rhs_oop = Systems.rhs_oop(sys) + Test.@test rhs_oop isa Function + + # Test OOP call with vector + u = [1.0, 2.0, 3.0, 4.0] # x = [1, 2], p = [3, 4] + p = Common.ODEParameters(nothing) + du = rhs_oop(u, p, 0.0) + + # ∂H/∂x = x = [1, 2], ∂H/∂p = p = [3, 4], so du = vcat(∂p, -∂x) = [3, 4, -1, -2] + Test.@test du == [3.0, 4.0, -1.0, -2.0] + end + + # ==================================================================== + # UNIT TESTS - build_rhs_augmented + # ==================================================================== + + Test.@testset "build_rhs_augmented" begin + h = Data.Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2) + 0.5 * v^2; is_autonomous=true, is_variable=false) + backend = FakeADBackend() + sys = Systems.HamiltonianSystem(h, backend; state_dimension=2) + + # Vector case (n_x=2, n_v=1) + rhs_aug = Systems.build_rhs_augmented(sys, 2, 1) + Test.@test rhs_aug isa Function + + u = [1.0, 2.0, 3.0, 4.0, 0.5] # x = [1, 2], p = [3, 4], pv = [0.5] + du = zeros(5) + p = Common.ODEParameters(0.5) # variable is 0.5 (matches pv in u) + rhs_aug(du, u, p, 0.0) + + # ∂H/∂x = x = [1, 2], ∂H/∂p = p = [3, 4], ∂H/∂v = v = 0.5 + # du = [∂p, -∂x, -∂v] = [3, 4, -1, -2, -0.5] + Test.@test du[1:2] == [3.0, 4.0] + Test.@test du[3:4] == [-1.0, -2.0] + Test.@test du[5] == -0.5 + + # Matrix compatible case (10×3, v matrice 1×3) + u_mat = [1.0 2.0 3.0; 4.0 5.0 6.0; 7.0 8.0 9.0; 10.0 11.0 12.0; 0.5 0.6 0.7] + du_mat = zeros(5, 3) + p_mat = Common.ODEParameters([0.5 0.6 0.7]) # v as matrix with 3 columns + rhs_aug(du_mat, u_mat, p_mat, 0.0) + Test.@test size(du_mat, 2) == 3 # no error, compatible + + # Matrix incompatible case (10×3, v matrice 1×2) → PreconditionError + u_mat2 = [1.0 2.0; 4.0 5.0; 7.0 8.0; 10.0 11.0; 0.5 0.6] + du_mat2 = zeros(5, 2) + p_mat2 = Common.ODEParameters([0.5 0.6 0.7]) # v has 3 columns, u has 2 + Test.@test_throws Exceptions.PreconditionError rhs_aug(du_mat2, u_mat2, p_mat2, 0.0) + + # u Matrix, v Vector (no-op check) + u_mat3 = [1.0 2.0; 4.0 5.0; 7.0 8.0; 10.0 11.0; 0.5 0.6] + du_mat3 = zeros(5, 2) + p_vec = Common.ODEParameters(0.5) # v as scalar (not a matrix) + rhs_aug(du_mat3, u_mat3, p_vec, 0.0) # no error, no-op check + end + + # ==================================================================== + # UNIT TESTS - _aug_split + # ==================================================================== + + Test.@testset "_aug_split" begin + # Vector case + u_vec = [1.0, 2.0, 3.0, 4.0, 0.5] + x, p, pv = Systems._aug_split(u_vec, 2, 1) + Test.@test x == [1.0, 2.0] + Test.@test p == [3.0, 4.0] + Test.@test pv == [0.5] + + # Matrix case + u_mat = [1.0 2.0; 3.0 4.0; 5.0 6.0; 7.0 8.0; 0.5 0.6] + x_mat, p_mat, pv_mat = Systems._aug_split(u_mat, 2, 1) + Test.@test x_mat == [1.0 2.0; 3.0 4.0] + Test.@test p_mat == [5.0 6.0; 7.0 8.0] + Test.@test pv_mat == [0.5 0.6] + end + + # ==================================================================== + # UNIT TESTS - _check_aug_batch_compat + # ==================================================================== + + Test.@testset "_check_aug_batch_compat" begin + # Compatible matrices + u = [1.0 2.0; 3.0 4.0] + v = [0.5 0.6] + Test.@test Systems._check_aug_batch_compat(u, v) === nothing # no error + + # Incompatible matrices + u2 = [1.0 2.0; 3.0 4.0] + v2 = [0.5 0.6 0.7] + Test.@test_throws Exceptions.PreconditionError Systems._check_aug_batch_compat(u2, v2) + + # Non-matrix cases (no-op) + u3 = [1.0, 2.0, 3.0] + v3 = [0.5] + Test.@test Systems._check_aug_batch_compat(u3, v3) === nothing + + u4 = [1.0 2.0; 3.0 4.0] + v4 = 0.5 + Test.@test Systems._check_aug_batch_compat(u4, v4) === nothing + end + + # ==================================================================== + # UNIT TESTS - build_system overloads + # ==================================================================== + + Test.@testset "build_system" begin + h = Data.Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2); is_autonomous=true, is_variable=false) + backend = FakeADBackend() + + # Without state_dimension + sys1 = Systems.build_system(h, backend) + Test.@test sys1 isa Systems.HamiltonianSystem + Test.@test Systems.state_dimension(sys1) === nothing + + # With state_dimension + sys2 = Systems.build_system(h, backend; state_dimension=2) + Test.@test sys2 isa Systems.HamiltonianSystem + Test.@test Systems.state_dimension(sys2) == 2 + end + + # ==================================================================== + # UNIT TESTS - Type Stability + # ==================================================================== + + Test.@testset "Type Stability" begin + h = Data.Hamiltonian((t, x, p, v) -> 0.5 * sum(x.^2) + sum(p.^2); is_autonomous=true, is_variable=false) + backend = FakeADBackend() + + sys1 = Systems.HamiltonianSystem(h, backend) + Test.@test Test.@inferred Systems.state_dimension(sys1) === nothing + Test.@test Test.@inferred Systems.ad_trait(sys1) === Common.WithAD + + sys2 = Systems.HamiltonianSystem(h, backend; state_dimension=2) + Test.@test Test.@inferred Systems.state_dimension(sys2) == 2 + Test.@test Test.@inferred Systems.ad_trait(sys2) === Common.WithAD + end + + # ==================================================================== + # REGRESSION TEST - HamiltonianVectorFieldSystem unaffected + # ==================================================================== + + Test.@testset "Regression: HamiltonianVectorFieldSystem" begin + hvf = Data.HamiltonianVectorField((x, p) -> (x, -p); is_autonomous=true, is_variable=false) + sys = Systems.HamiltonianVectorFieldSystem(hvf) + Test.@test sys isa Systems.HamiltonianVectorFieldSystem + Test.@test sys isa Systems.AbstractHamiltonianSystem + Test.@test Systems.ad_trait(sys) === Common.WithoutAD + end + end +end + +end # module + +test_hamiltonian_system() = TestHamiltonianSystem.test_hamiltonian_system() diff --git a/test/suite/systems/test_hamiltonian_vector_field_system.jl b/test/suite/systems/test_hamiltonian_vector_field_system.jl index f40c5998..2bf2ff97 100644 --- a/test/suite/systems/test_hamiltonian_vector_field_system.jl +++ b/test/suite/systems/test_hamiltonian_vector_field_system.jl @@ -25,11 +25,22 @@ function test_hamiltonian_vector_field_system() Test.@test sys1 isa Systems.HamiltonianVectorFieldSystem Test.@test sys1 isa Systems.AbstractHamiltonianSystem Test.@test Systems.state_dimension(sys1) === nothing - + # From HamiltonianVectorField with dimension sys2 = Systems.HamiltonianVectorFieldSystem(hvf; state_dimension=3) Test.@test sys2 isa Systems.HamiltonianVectorFieldSystem Test.@test Systems.state_dimension(sys2) == 3 + + # Hierarchy check: supertype with WithoutAD + Test.@test sys1 isa Systems.AbstractHamiltonianSystem{Common.Autonomous, Common.Fixed, Common.WithoutAD} + end + + Test.@testset "ad_trait" begin + hvf = Data.HamiltonianVectorField((x, p) -> (x, -p); is_autonomous=true, is_variable=false) + sys = Systems.HamiltonianVectorFieldSystem(hvf) + + Test.@test Systems.ad_trait(sys) === Common.WithoutAD + Test.@test Test.@inferred Systems.ad_trait(sys) === Common.WithoutAD end # ====================================================================