complete v0.2.0 upgrade

## Breaking changes

- Refactored the generation method of the RuntimeGenerate function:
    - Added support for ComponentArray type as parameter input, eliminating the need to convert it to Vector type.
    - Enabled broadcasting calculations between vectors and scalars in the RuntimeGenerate function, improving computational efficiency and gradient computation efficiency.
- Modified the construction method of Route, now using rfluxes, dfluxes, and proj_func to build the Route type.
- Revised the computational logic of Route, Bucket, and Flux, providing corresponding functions for both three-dimensional and two-dimensional data inputs, avoiding some complex internal data integration operations.
- Added a macro construction method (under test).

Author: chooron

Manifest.toml

@@ -4,7 +4,9 @@ authors = ["jingx <50790703+chooron@users.noreply.github.com>"]
 version = "0.2.0"
 
 [deps]
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 HydroModelTools = "31f4d4b3-d71d-422d-8932-b4ab24c6e7e3"

Project.toml

@@ -29,6 +29,7 @@
     - 使用@generated的函数时,只能获取到类型,具体的实例属性是获取不到的
 - [ ] AbstractBucket分成HydroBucket和MultiBucket
 - [ ] 参数名称与输出输入名称不能冲突
+- [ ] Meta这个存储的内容得更改了
 
 
 julia语言中,我想使用@generated生成struct函数,但是这个生成函数需要的参数需要struct的属性值,这个属性值是术语Num,因此无法作为struct的类型参数,请问应该如何解决这个问题呢

TODO.md

@@ -0,0 +1,8 @@
+## Breaking changes
+
+- Refactored the generation method of the RuntimeGenerate function:
+    - Added support for ComponentArray type as parameter input, eliminating the need to convert it to Vector type.
+    - Enabled broadcasting calculations between vectors and scalars in the RuntimeGenerate function, improving computational efficiency and gradient computation efficiency.
+- Modified the construction method of Route, now using rfluxes, dfluxes, and proj_func to build the Route type.
+- Revised the computational logic of Route, Bucket, and Flux, providing corresponding functions for both three-dimensional and two-dimensional data inputs, avoiding some complex internal data integration operations.
+- Added a macro construction method (under test).

release_note.txt

@@ -24,7 +24,6 @@ using SymbolicUtils.Code
 import SymbolicUtils: BasicSymbolic, Sym, Term, iscall, operation, arguments, issym, symtype, sorted_arguments
 @reexport using ModelingToolkit: @variables, @parameters
 using ModelingToolkit: isparameter
-using ModelingToolkit: t_nounits as t
 # graph compute
 using Graphs
 
@@ -38,9 +37,6 @@ using NNlib
 
 ## Abstract Component Types
 abstract type AbstractComponent end
-abstract type AbstractIOAdapter end
-abstract type AbstractHydroWrapper <: AbstractComponent end
-abstract type AbstractNeuralWrapper <: AbstractComponent end
 
 abstract type AbstractFlux <: AbstractComponent end
 abstract type AbstractHydroFlux <: AbstractFlux end
@@ -54,16 +50,15 @@ abstract type AbstractRoute <: AbstractElement end
 abstract type AbstractHydroRoute <: AbstractRoute end
 abstract type AbstractModel <: AbstractComponent end
 
-export AbstractComponent, AbstractHydroWrapper, AbstractNeuralWrapper
-export AbstractFlux, AbstractHydroFlux, AbstractNeuralFlux, AbstractStateFlux
+export AbstractComponent, AbstractFlux, AbstractHydroFlux, AbstractNeuralFlux, AbstractStateFlux
 export AbstractElement, AbstractBucket, AbstractHydrograph, AbstractRoute, AbstractHydroRoute, AbstractModel
 
 # utils
 include("utils/expression.jl")
 include("utils/attribute.jl")
+include("utils/tools.jl")
 include("utils/display.jl")
 include("utils/build.jl")
-include("utils/sort.jl")
 include("utils/check.jl")
 #! A discrete ODE solver, if want to use more efficient solver, please import HydroModelTools.jl
 include("utils/solver.jl")

src/HydroModels.jl

@@ -1,40 +1,60 @@
 """
-	HydroBucket(; funcs::Vector{<:AbstractHydroFlux}, dfuncs::Vector{<:AbstractStateFlux}=StateFlux[], name::Union{Symbol,Nothing}=nothing, sort_funcs::Bool=false)
+    HydroBucket(; 
+        funcs::Vector{<:AbstractHydroFlux}, 
+        dfuncs::Vector{<:AbstractStateFlux}=StateFlux[], 
+        name::Union{Symbol,Nothing}=nothing, 
+        sort_funcs::Bool=false
+    )
 
-Represents a hydrological bucket model component.
+Represents a hydrological bucket model component that handles both single-node and multi-node computations.
 
 # Arguments
-- `funcs::Vector{<:AbstractHydroFlux}`: A vector of flux functions that describe the hydrological processes.
-- `dfuncs::Vector{<:AbstractStateFlux}`: A vector of state derivative functions (default is an empty vector of StateFlux).
-- `name::Union{Symbol,Nothing}`: Optional name for the bucket. If not provided, a name will be automatically generated from state variable names.
-- `sort_funcs::Bool`: Whether to sort the flux functions (default is false).
+- `funcs::Vector{<:AbstractHydroFlux}`: Vector of flux functions describing hydrological processes
+- `dfuncs::Vector{<:AbstractStateFlux}`: Vector of state derivative functions for ODE calculations (default: empty)
+- `name::Union{Symbol,Nothing}`: Optional bucket identifier. Defaults to auto-generated name from state variables
+- `sort_funcs::Bool}`: Whether to topologically sort flux functions based on dependencies (default: false)
 
 # Fields
-- `funcs::Vector{<:AbstractHydroFlux}`: Vector of flux functions describing hydrological processes.
-- `dfuncs::Vector{<:AbstractStateFlux}`: Vector of state derivative functions for ODE calculations.
-- `flux_func::Function`: Combined function for calculating all hydrological fluxes.
-- `ode_func::Union{Nothing,Function}`: Function for ordinary differential equations (ODE) calculations, or nothing if no ODE calculations are needed.
-- `meta::HydroMeta`: Contains metadata about the bucket, including input, output, state, parameter, and neural network names.
+- `fluxes::Vector{<:AbstractHydroFlux}`: Vector of flux functions describing hydrological processes
+- `dfluxes::Vector{<:AbstractStateFlux}`: Vector of state derivative functions for ODE calculations
+- `flux_func::Function`: Generated function for single-node flux calculations
+- `multi_flux_func::Function`: Generated function for multi-node parallel flux calculations
+- `ode_func::Union{Nothing,Function}`: Generated function for single-node ODE calculations
+- `multi_ode_func::Union{Nothing,Function}`: Generated function for multi-node parallel ODE calculations
+- `meta::ComponentVector`: Metadata containing model structure information
 
 # Description
-HydroBucket is a structure that encapsulates the behavior of a hydrological bucket model. 
-It combines multiple flux functions and state derivative functions to model water movement 
-and storage within a hydrological unit.
-
-The structure automatically extracts relevant information from the provided functions to 
-populate the metadata, which includes names of:
-- Inputs: Variables that drive the model
-- Outputs: Variables produced by the model
-- States: Internal model states that evolve over time
-- Parameters: Model parameters that control behavior
-- Neural Networks: Any neural network components (if applicable)
-
-The `flux_func` and `ode_func` are constructed based on the provided `funcs` and `dfuncs`, 
-enabling efficient calculation of fluxes and state changes over time.
-
-This structure is particularly useful for building complex hydrological models by combining 
-multiple HydroBucket instances to represent different components of a water system.
-
+HydroBucket is a type-stable implementation of a hydrological bucket model that supports both 
+single-node and distributed (multi-node) computations. It automatically generates optimized 
+functions for flux and ODE calculations based on the provided process functions.
+
+## Model Structure
+The bucket model consists of:
+- Process functions (`fluxes`): Define water movement between storages
+- State derivatives (`dfluxes`): Define rate of change for state variables
+- Generated functions: Optimized implementations for both single and multi-node calculations
+
+## Metadata Components
+The `meta` field tracks:
+- `inputs`: External forcing variables (e.g., precipitation, temperature)
+- `outputs`: Model-generated variables (e.g., runoff, evaporation)
+- `states`: Internal storage variables (e.g., soil moisture, groundwater)
+- `params`: Model parameters controlling process behavior
+- `nn_vars`: Neural network components (if any) for hybrid modeling
+
+## Performance Features
+- Type-stable computations for both single and multi-node cases
+- Efficient broadcasting operations for vectorized calculations
+- Automatic function generation with optimized broadcasting
+- Support for ComponentArray parameters for structured data handling
+
+## Usage Notes
+1. For single-node simulations: Use `flux_func` and `ode_func`
+2. For multi-node simulations: Use `multi_flux_func` and `multi_ode_func`
+3. Parameters should be provided as ComponentVector for type stability
+4. Broadcasting operations are automatically handled for multi-node cases
+
+See also: [`AbstractHydroFlux`](@ref), [`AbstractStateFlux`](@ref), [`ComponentVector`](@ref)
 """
 struct HydroBucket{S} <: AbstractBucket
     "Name of the bucket"
@@ -43,13 +63,13 @@ struct HydroBucket{S} <: AbstractBucket
     fluxes::Vector{<:AbstractHydroFlux}
     "Vector of state derivative functions for ODE calculations."
     dfluxes::Vector{<:AbstractStateFlux}
-    "Generated function for calculating all hydrological fluxes."
+    "Generated function for calculating all hydrological fluxes. (Supports single-node data)"
     flux_func::Function
-    "Generated function for calculating all hydrological fluxes."
+    "Generated function for calculating all hydrological fluxes. (Supports multi-nodes data)"
     multi_flux_func::Function
-    "Generated function for ordinary differential equations (ODE) calculations, or nothing if no ODE calculations are needed."
+    "Generated function for ordinary differential equations (ODE) calculations, or nothing if no ODE calculations are needed. (Supports single-node data)"
     ode_func::Union{Nothing,Function}
-    "Generated function for ordinary differential equations (ODE) calculations, or nothing if no ODE calculations are needed."
+    "Generated function for ordinary differential equations (ODE) calculations, or nothing if no ODE calculations are needed. (Supports multi-nodes data)"
     multi_ode_func::Union{Nothing,Function}
     "Metadata about the bucket, including input, output, state, parameter, and neural network names."
     meta::ComponentVector
@@ -118,10 +138,8 @@ The input dimensions must match the number of input variables defined in the mod
 Required parameters and initial states must be present in the pas argument.
 """
 function (ele::HydroBucket{true})(
-    input::AbstractArray{T,2},
-    pas::ComponentVector;
-    config::NamedTuple=NamedTuple(),
-    kwargs...,
+    input::AbstractArray{T,2}, pas::ComponentVector;
+    config::NamedTuple=NamedTuple(), kwargs...
 ) where {T}
     #* get kwargs
     solver = get(config, :solver, ManualSolver{true}())
@@ -139,15 +157,13 @@ function (ele::HydroBucket{true})(
 end
 
 (ele::HydroBucket{false})(input::AbstractArray{T,2}, pas::ComponentVector; kwargs...) where {T} = begin
-    flux_output = ele.flux_func(input, nothing, pas)
+    flux_output = ele.flux_func(eachslice(input, dims=1), nothing, pas)
     permutedims(reduce(hcat, flux_output))
 end
 
 function (ele::HydroBucket{true})(
-    input::AbstractArray{T,3},
-    pas::ComponentVector;
-    config::NamedTuple=NamedTuple(),
-    kwargs...,
+    input::AbstractArray{T,3}, pas::ComponentVector;
+    config::NamedTuple=NamedTuple(), kwargs...
 ) where {T}
     input_dims, num_nodes, time_len = size(input)
 
@@ -159,15 +175,11 @@ function (ele::HydroBucket{true})(
     timeidx = get(config, :timeidx, collect(1:size(input, 3)))
 
     #* prepare states parameters and nns
-    params = view(pas, :params)
-    nn_params = isempty(get_nn_vars(ele)) ? ones(eltype(pas), num_nodes) : view(pas, :nns)
-    expand_params = ComponentVector(NamedTuple{Tuple(get_param_names(ele))}([params[p][ptyidx] for p in get_param_names(ele)]))
-    new_pas = ComponentVector(params=expand_params, nns=nn_params)
-    initstates_mat = view(reshape(Vector(view(pas, :initstates)), num_nodes, :)', :, styidx)
+    new_pas = expand_component_params(pas, ptyidx)
+    initstates_mat = expand_component_initstates(pas, styidx)
 
     #* prepare input function
-    input_reshape = reshape(input, input_dims * num_nodes, time_len)
-    itpfuncs = interp(input_reshape, timeidx)
+    itpfuncs = interp(reshape(input, input_dims * num_nodes, time_len), timeidx)
     solved_states = solver(
         (u, p, t) -> begin
             tmp_input = reshape(itpfuncs(t), input_dims, num_nodes)
@@ -182,16 +194,11 @@ function (ele::HydroBucket{true})(
 end
 
 function (ele::HydroBucket{false})(
-    input::AbstractArray{T,3},
-    pas::ComponentVector;
-    config::NamedTuple=NamedTuple(),
-    kwargs...,
+    input::AbstractArray{T,3}, pas::ComponentVector;
+    config::NamedTuple=NamedTuple(), kwargs...,
 ) where {T}
     ptyidx = get(config, :ptyidx, 1:size(input, 2))
-    params = view(pas, :params)
-    nn_params = isempty(get_nn_vars(ele)) ? ones(eltype(pas), size(input, 2)) : view(pas, :nns)
-    expand_params = ComponentVector(NamedTuple{Tuple(get_param_names(ele))}([params[p][ptyidx] for p in get_param_names(ele)]))
-    new_pas = ComponentVector(params=expand_params, nns=nn_params)
+    new_pas = expand_component_params(pas, ptyidx)
     #* run other functions
     output = ele.flux_func(eachslice(input, dims=1), nothing, new_pas)
     permutedims(reduce((m1, m2) -> cat(m1, m2, dims=3), output), (3, 1, 2))

src/bucket.jl

@@ -139,6 +139,8 @@ where traditional equations may be insufficient or unknown.
 struct NeuralFlux <: AbstractNeuralFlux
     "Name of the flux"
     name::Symbol
+    "chain of the neural network"
+    chain::LuxCore.AbstractLuxLayer
     "Compiled function that calculates the flux using the neural network"
     func::Function
     "Metadata about the flux, including input, output, and neural network parameter names"
@@ -154,7 +156,6 @@ struct NeuralFlux <: AbstractNeuralFlux
         chain_name::Union{Symbol,Nothing}=nothing,
     ) where {T<:Num}
         #* Check chain name
-        @assert chain.name isa Symbol "Neural network chain should have a name with Symbol type"
         chain_name = chain_name === nothing ? chain.name : chain_name
 
         ps, st = Lux.setup(StableRNG(42), chain)
@@ -165,8 +166,8 @@ struct NeuralFlux <: AbstractNeuralFlux
 
         meta = ComponentVector(inputs=inputs, outputs=outputs, nns=NamedTuple{Tuple([chain_name])}([chain_params]))
         nninfos = (inputs=nn_input_name, outputs=nn_output_name, nns=chain_name)
-        flux_name = isnothing(name) ? Symbol("##neural_flux#", meta) : name
-        new(flux_name, nn_func, meta, nninfos)
+        flux_name = isnothing(name) ? Symbol("##neural_flux#", hash(meta)) : name
+        new(flux_name, chain, nn_func, meta, nninfos)
     end
 
     #* construct neural flux with input fluxes and output fluxes

src/flux.jl

@@ -66,7 +66,7 @@ struct HydroRoute <: AbstractHydroRoute
         #* define the route name
         route_name = isnothing(name) ? Symbol("##route#", hash(meta)) : name
         #* build the route function
-        multi_flux_func, multi_ode_func = build_route_func(rfluxes, dfluxes, proj_func, meta)
+        multi_flux_func, multi_ode_func = build_route_func(rfluxes, dfluxes, meta)
         return new(route_name, rfluxes, multi_flux_func, multi_ode_func, proj_func, meta)
     end
 end
@@ -223,11 +223,8 @@ function (route::HydroRoute)(
     timeidx = get(config, :timeidx, collect(1:time_len))
 
     #* prepare states parameters and nns
-    params = view(pas, :params)
-    expand_params = ComponentVector(NamedTuple{Tuple(get_param_names(route))}([params[p][ptyidx] for p in get_param_names(route)]))
-    nn_params = isempty(get_nn_vars(route)) ? Vector{eltype(pas)}[] : view(pas, :nns)
-    new_pas = ComponentVector(params=expand_params) # , nns=nn_params
-    initstates_mat = view(reshape(Vector(view(pas, :initstates)), num_nodes, :)', :, styidx)
+    new_pas = expand_component_params(pas, ptyidx)
+    initstates_mat = expand_component_initstates(pas, styidx)
 
     #* prepare input function
     input_reshape = reshape(input, input_dims * num_nodes, time_len)