|
| 1 | +""" |
| 2 | +$(TYPEDEF) |
| 3 | +
|
| 4 | +A constraint system of equations. |
| 5 | +
|
| 6 | +# Fields |
| 7 | +$(FIELDS) |
| 8 | +
|
| 9 | +# Examples |
| 10 | +
|
| 11 | +```julia |
| 12 | +@variables x y z |
| 13 | +@parameters a b c |
| 14 | +
|
| 15 | +cstr = [0 ~ a*(y-x), |
| 16 | + 0 ~ x*(b-z)-y, |
| 17 | + 0 ~ x*y - c*z |
| 18 | + x^2 + y^2 ≲ 1] |
| 19 | +@named ns = ConstraintsSystem(cstr, [x,y,z],[a,b,c]) |
| 20 | +``` |
| 21 | +""" |
| 22 | +struct ConstraintsSystem <: AbstractTimeIndependentSystem |
| 23 | + """ |
| 24 | + tag: a tag for the system. If two system have the same tag, then they are |
| 25 | + structurally identical. |
| 26 | + """ |
| 27 | + tag::UInt |
| 28 | + """Vector of equations defining the system.""" |
| 29 | + constraints::Vector{Union{Equation, Inequality}} |
| 30 | + """Unknown variables.""" |
| 31 | + states::Vector |
| 32 | + """Parameters.""" |
| 33 | + ps::Vector |
| 34 | + """Array variables.""" |
| 35 | + var_to_name::Any |
| 36 | + """Observed variables.""" |
| 37 | + observed::Vector{Equation} |
| 38 | + """ |
| 39 | + Jacobian matrix. Note: this field will not be defined until |
| 40 | + [`calculate_jacobian`](@ref) is called on the system. |
| 41 | + """ |
| 42 | + jac::RefValue{Any} |
| 43 | + """ |
| 44 | + Name: the name of the system. These are required to have unique names. |
| 45 | + """ |
| 46 | + name::Symbol |
| 47 | + """ |
| 48 | + systems: The internal systems |
| 49 | + """ |
| 50 | + systems::Vector{ConstraintsSystem} |
| 51 | + """ |
| 52 | + defaults: The default values to use when initial conditions and/or |
| 53 | + parameters are not supplied in `ODEProblem`. |
| 54 | + """ |
| 55 | + defaults::Dict |
| 56 | + """ |
| 57 | + type: type of the system |
| 58 | + """ |
| 59 | + connector_type::Any |
| 60 | + """ |
| 61 | + metadata: metadata for the system, to be used by downstream packages. |
| 62 | + """ |
| 63 | + metadata::Any |
| 64 | + """ |
| 65 | + tearing_state: cache for intermediate tearing state |
| 66 | + """ |
| 67 | + tearing_state::Any |
| 68 | + """ |
| 69 | + substitutions: substitutions generated by tearing. |
| 70 | + """ |
| 71 | + substitutions::Any |
| 72 | + |
| 73 | + function ConstraintsSystem(tag, constraints, states, ps, var_to_name, observed, jac, |
| 74 | + name, |
| 75 | + systems, |
| 76 | + defaults, connector_type, metadata = nothing, |
| 77 | + tearing_state = nothing, substitutions = nothing; |
| 78 | + checks::Union{Bool, Int} = true) |
| 79 | + if checks == true || (checks & CheckUnits) > 0 |
| 80 | + all_dimensionless([states; ps]) || check_units(constraints) |
| 81 | + end |
| 82 | + new(tag, constraints, states, ps, var_to_name, observed, jac, name, systems, |
| 83 | + defaults, |
| 84 | + connector_type, metadata, tearing_state, substitutions) |
| 85 | + end |
| 86 | +end |
| 87 | + |
| 88 | +equations(sys::ConstraintsSystem) = constraints(sys) # needed for Base.show |
| 89 | + |
| 90 | +function ConstraintsSystem(constraints, states, ps; |
| 91 | + observed = [], |
| 92 | + name = nothing, |
| 93 | + default_u0 = Dict(), |
| 94 | + default_p = Dict(), |
| 95 | + defaults = _merge(Dict(default_u0), Dict(default_p)), |
| 96 | + systems = ConstraintsSystem[], |
| 97 | + connector_type = nothing, |
| 98 | + continuous_events = nothing, # this argument is only required for ODESystems, but is added here for the constructor to accept it without error |
| 99 | + discrete_events = nothing, # this argument is only required for ODESystems, but is added here for the constructor to accept it without error |
| 100 | + checks = true, |
| 101 | + metadata = nothing) |
| 102 | + continuous_events === nothing || isempty(continuous_events) || |
| 103 | + throw(ArgumentError("ConstraintsSystem does not accept `continuous_events`, you provided $continuous_events")) |
| 104 | + discrete_events === nothing || isempty(discrete_events) || |
| 105 | + throw(ArgumentError("ConstraintsSystem does not accept `discrete_events`, you provided $discrete_events")) |
| 106 | + |
| 107 | + name === nothing && |
| 108 | + throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro")) |
| 109 | + |
| 110 | + cstr = value.(Symbolics.canonical_form.(scalarize(constraints))) |
| 111 | + states′ = value.(scalarize(states)) |
| 112 | + ps′ = value.(scalarize(ps)) |
| 113 | + |
| 114 | + if !(isempty(default_u0) && isempty(default_p)) |
| 115 | + Base.depwarn("`default_u0` and `default_p` are deprecated. Use `defaults` instead.", |
| 116 | + :ConstraintsSystem, force = true) |
| 117 | + end |
| 118 | + sysnames = nameof.(systems) |
| 119 | + if length(unique(sysnames)) != length(sysnames) |
| 120 | + throw(ArgumentError("System names must be unique.")) |
| 121 | + end |
| 122 | + |
| 123 | + jac = RefValue{Any}(EMPTY_JAC) |
| 124 | + defaults = todict(defaults) |
| 125 | + defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults)) |
| 126 | + |
| 127 | + var_to_name = Dict() |
| 128 | + process_variables!(var_to_name, defaults, states′) |
| 129 | + process_variables!(var_to_name, defaults, ps′) |
| 130 | + isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed)) |
| 131 | + |
| 132 | + ConstraintsSystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)), |
| 133 | + cstr, states, ps, var_to_name, observed, jac, name, systems, |
| 134 | + defaults, |
| 135 | + connector_type, metadata, checks = checks) |
| 136 | +end |
| 137 | + |
| 138 | +function calculate_jacobian(sys::ConstraintsSystem; sparse = false, simplify = false) |
| 139 | + cache = get_jac(sys)[] |
| 140 | + if cache isa Tuple && cache[2] == (sparse, simplify) |
| 141 | + return cache[1] |
| 142 | + end |
| 143 | + |
| 144 | + lhss = generate_canonical_form_lhss(sys) |
| 145 | + vals = [dv for dv in states(sys)] |
| 146 | + if sparse |
| 147 | + jac = sparsejacobian(lhss, vals, simplify = simplify) |
| 148 | + else |
| 149 | + jac = jacobian(lhss, vals, simplify = simplify) |
| 150 | + end |
| 151 | + get_jac(sys)[] = jac, (sparse, simplify) |
| 152 | + return jac |
| 153 | +end |
| 154 | + |
| 155 | +function generate_jacobian(sys::ConstraintsSystem, vs = states(sys), ps = parameters(sys); |
| 156 | + sparse = false, simplify = false, kwargs...) |
| 157 | + jac = calculate_jacobian(sys, sparse = sparse, simplify = simplify) |
| 158 | + return build_function(jac, vs, ps; kwargs...) |
| 159 | +end |
| 160 | + |
| 161 | +function calculate_hessian(sys::ConstraintsSystem; sparse = false, simplify = false) |
| 162 | + lhss = generate_canonical_form_lhss(sys) |
| 163 | + vals = [dv for dv in states(sys)] |
| 164 | + if sparse |
| 165 | + hess = [sparsehessian(lhs, vals, simplify = simplify) for lhs in lhss] |
| 166 | + else |
| 167 | + hess = [hessian(lhs, vals, simplify = simplify) for lhs in lhss] |
| 168 | + end |
| 169 | + return hess |
| 170 | +end |
| 171 | + |
| 172 | +function generate_hessian(sys::ConstraintsSystem, vs = states(sys), ps = parameters(sys); |
| 173 | + sparse = false, simplify = false, kwargs...) |
| 174 | + hess = calculate_hessian(sys, sparse = sparse, simplify = simplify) |
| 175 | + return build_function(hess, vs, ps; kwargs...) |
| 176 | +end |
| 177 | + |
| 178 | +function generate_function(sys::ConstraintsSystem, dvs = states(sys), ps = parameters(sys); |
| 179 | + kwargs...) |
| 180 | + lhss = generate_canonical_form_lhss(sys) |
| 181 | + pre, sol_states = get_substitutions_and_solved_states(sys) |
| 182 | + |
| 183 | + func = build_function(lhss, value.(dvs), value.(ps); postprocess_fbody = pre, |
| 184 | + states = sol_states, kwargs...) |
| 185 | + |
| 186 | + cstr = constraints(sys) |
| 187 | + lcons = fill(-Inf, length(cstr)) |
| 188 | + ucons = zeros(length(cstr)) |
| 189 | + lcons[findall(Base.Fix2(isa, Equation), cstr)] .= 0.0 |
| 190 | + |
| 191 | + return func, lcons, ucons |
| 192 | +end |
| 193 | + |
| 194 | +function jacobian_sparsity(sys::ConstraintsSystem) |
| 195 | + lhss = generate_canonical_form_lhss(sys) |
| 196 | + jacobian_sparsity(lhss, states(sys)) |
| 197 | +end |
| 198 | + |
| 199 | +function hessian_sparsity(sys::ConstraintsSystem) |
| 200 | + lhss = generate_canonical_form_lhss(sys) |
| 201 | + [hessian_sparsity(eq, states(sys)) for eq in lhss] |
| 202 | +end |
| 203 | + |
| 204 | +""" |
| 205 | +Convert the system of equalities and inequalities into a canonical form: |
| 206 | +h(x) = 0 |
| 207 | +g(x) <= 0 |
| 208 | +""" |
| 209 | +function generate_canonical_form_lhss(sys) |
| 210 | + lhss = [Symbolics.canonical_form(eq).lhs for eq in constraints(sys)] |
| 211 | +end |
0 commit comments