Add HamiltonianSystem with AD support and AugmentedHamiltonianTrait

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

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)
 # =============================================================================

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
 # =============================================================================

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)

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}

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

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

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
+