@@ -2583,27 +2583,15 @@ def has_event_assignments(self) -> bool:
25832583 """
25842584 return any (event .updates_state for event in self ._events )
25852585
2586- def toposort_expressions (self ) -> dict [sp .Symbol , sp .Expr ]:
2586+ def toposort_expressions (
2587+ self , reorder : bool = True
2588+ ) -> dict [sp .Symbol , sp .Expr ]:
25872589 """
25882590 Sort expressions in topological order.
25892591
25902592 :return:
25912593 dict of expression symbols to expressions in topological order
25922594 """
2593- # ensure no symbols or equations that depend on `w` have been generated
2594- # yet, otherwise the re-ordering might break dependencies
2595- if (
2596- generated := set (self ._syms )
2597- | set (self ._eqs )
2598- | set (self ._sparsesyms )
2599- | set (self ._sparseeqs )
2600- ) - {"w" , "p" , "k" , "x" , "x_rdata" }:
2601- raise AssertionError (
2602- "This function must be called before computing any "
2603- "derivatives. The following symbols/equations are already "
2604- f"generated: { generated } "
2605- )
2606-
26072595 # NOTE: elsewhere, conservations law expressions are expected to
26082596 # occur before any other w expressions, so we must maintain their
26092597 # position.
@@ -2622,6 +2610,23 @@ def toposort_expressions(self) -> dict[sp.Symbol, sp.Expr]:
26222610 for e in self .expressions ()[: self .num_cons_law ()]
26232611 } | w_toposorted
26242612
2613+ if not reorder :
2614+ return w_toposorted
2615+
2616+ # ensure no symbols or equations that depend on `w` have been generated
2617+ # yet, otherwise the re-ordering might break dependencies
2618+ if (
2619+ generated := set (self ._syms )
2620+ | set (self ._eqs )
2621+ | set (self ._sparsesyms )
2622+ | set (self ._sparseeqs )
2623+ ) - {"w" , "p" , "k" , "x" , "x_rdata" }:
2624+ raise AssertionError (
2625+ "This function must be called before computing any "
2626+ "derivatives. The following symbols/equations are already "
2627+ f"generated: { generated } "
2628+ )
2629+
26252630 old_syms = tuple (e .get_sym () for e in self .expressions ())
26262631 topo_expr_syms = tuple (w_toposorted )
26272632 new_order = [old_syms .index (s ) for s in topo_expr_syms ]
0 commit comments