Skip to content

Commit 45e4843

Browse files
feat: add flatten_equation utility
Split out core of `flatten_equations`
1 parent 8a14d18 commit 45e4843

File tree

1 file changed

+20
-9
lines changed

1 file changed

+20
-9
lines changed

lib/ModelingToolkitBase/src/utils.jl

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1191,6 +1191,25 @@ end
11911191

11921192
_eq_unordered(a, b) = isequal(a, b)
11931193

1194+
"""
1195+
$TYPEDSIGNATURES
1196+
1197+
Given an equation that may be an array equation, return a `Vector{Equation}` representing
1198+
its scalarized form.
1199+
"""
1200+
function flatten_equation(eq::Equation)::Vector{Equation}
1201+
if !SU.is_array_shape(SU.shape(eq.lhs))
1202+
return [eq]
1203+
end
1204+
lhs = vec(collect(eq.lhs)::Array{SymbolicT})::Vector{SymbolicT}
1205+
rhs = vec(collect(eq.rhs)::Array{SymbolicT})::Vector{SymbolicT}
1206+
result = Equation[]
1207+
for (l, r) in zip(lhs, rhs)
1208+
push!(result, l ~ r)
1209+
end
1210+
return result
1211+
end
1212+
11941213
"""
11951214
$(TYPEDSIGNATURES)
11961215
@@ -1200,15 +1219,7 @@ without scalarizing occurrences of array variables and return the new list of eq
12001219
function flatten_equations(eqs::Vector{Equation})
12011220
_eqs = Equation[]
12021221
for eq in eqs
1203-
if !SU.is_array_shape(SU.shape(eq.lhs))
1204-
push!(_eqs, eq)
1205-
continue
1206-
end
1207-
lhs = vec(collect(eq.lhs)::Array{SymbolicT})::Vector{SymbolicT}
1208-
rhs = vec(collect(eq.rhs)::Array{SymbolicT})::Vector{SymbolicT}
1209-
for (l, r) in zip(lhs, rhs)
1210-
push!(_eqs, l ~ r)
1211-
end
1222+
append!(_eqs, flatten_equation(eq))
12121223
end
12131224
return _eqs
12141225
end

0 commit comments

Comments
 (0)