From eaf6aced661a3447fbd0b90dffd3740a5ece4f92 Mon Sep 17 00:00:00 2001
From: Doug A <douglas.allan@keylogic.com>
Date: Wed, 26 Jun 2024 12:41:13 -0400
Subject: [PATCH 1/4] get_initial_condition_problem

---
 idaes/core/solvers/petsc.py            | 217 +++++++++++++++-----
 idaes/core/solvers/tests/test_petsc.py | 272 +++++++++++++++++++++----
 2 files changed, 397 insertions(+), 92 deletions(-)

diff --git a/idaes/core/solvers/petsc.py b/idaes/core/solvers/petsc.py
index 20252891e4..fdce3adf60 100644
--- a/idaes/core/solvers/petsc.py
+++ b/idaes/core/solvers/petsc.py
@@ -42,6 +42,7 @@
 import idaes
 import idaes.logger as idaeslog
 import idaes.config as icfg
+from idaes.core.util.model_statistics import degrees_of_freedom, number_activated_constraints
 
 
 # Importing a few things here so that they are cached
@@ -249,17 +250,19 @@ def _copy_time(time_vars, t_from, t_to):
 
 
 def find_discretization_equations(m, time):
-    """This returns a list of time discretization equations. Since we aren't
-    solving the whole time period simultaneously, we'll want to deactivate
-    these constraints.
+    """
+    This returns a list of continuity equations (if Lagrange-Legendre collocation is used)
+    and time discretization equations. Since we aren't solving the whole time period 
+    simultaneously, we'll want to deactivate these constraints.
 
     Args:
         m (Block): model or block to search for constraints
         time (ContinuousSet):
 
     Returns:
-        list of time discretization constraints
+        list of time continuity (if present) and discretization constraints
     """
+    tname = time.local_name
     disc_eqns = []
     for var in m.component_objects(pyo.Var):
         if isinstance(var, pyodae.DerivativeVar):
@@ -271,12 +274,40 @@ def find_discretization_equations(m, time):
                         "that are differentiated at least once with respect to time. Please reformulate your model so "
                         "it does not contain such a derivative (such as by introducing intermediate variables)."
                     )
-                parent = var.parent_block()
+                state_var = var.get_state_var()
+                parent_block = var.parent_block()
                 name = var.local_name + "_disc_eq"
-                disc_eq = getattr(parent, name)
-                disc_eqns.append(disc_eq)
+                disc_eqn = getattr(parent_block, name)
+                disc_eqns.append(disc_eqn)
+                try:
+                    cont_eqn = getattr(
+                        parent_block, state_var.local_name + "_" + tname + "_cont_eq"
+                    )
+                except AttributeError:
+                    cont_eqn = None
+                if cont_eqn is not None:
+                    disc_eqns.append(cont_eqn)
+
     return disc_eqns
 
+def get_discretization_constraint_data(m, time):
+    """
+    This returns a list of continuity equations (if Lagrange-Legendre collocation is used)
+    and time discretization equations, expanded into constraint_data form.
+
+    Args:
+        m (Block): model or block to search for constraints
+        time (ContinuousSet):
+
+    Returns:
+        list of time continuity (if present) and discretization constraint data
+    """
+    disc_eqns = find_discretization_equations(m, time)
+    con_data_list = []
+    for con in disc_eqns:
+        con_data_list += [con_data for con_data in con.values()]
+    return con_data_list
+
 
 def _set_dae_suffixes_from_variables(m, variables, deriv_diff_map):
     """Write suffixes used by the solver to identify different variable types
@@ -536,13 +567,11 @@ def petsc_dae_by_time_element(
     elif representative_time not in between:
         raise RuntimeError("representative_time is not element of between.")
 
-    regular_vars, time_vars = flatten_dae_components(
-        m, time, pyo.Var, active=True, indices=(representative_time,)
+    flattened_problem  = get_flattened_problem(
+        model=m,
+        time=time,
+        representative_time=representative_time
     )
-    regular_cons, time_cons = flatten_dae_components(
-        m, time, pyo.Constraint, active=True, indices=(representative_time,)
-    )
-    tdisc = find_discretization_equations(m, time)
 
     solver_dae = pyo.SolverFactory("petsc_ts", options=ts_options)
     save_trajectory = solver_dae.options.get("--ts_save_trajectory", 0)
@@ -554,7 +583,7 @@ def petsc_dae_by_time_element(
     if initial_variables is None:
         initial_variables = []
     if detect_initial:
-        rvset = ComponentSet(regular_vars)
+        rvset = ComponentSet(flattened_problem["timeless_variables"])
         ivset = ComponentSet(initial_variables)
         initial_variables = list(ivset | rvset)
 
@@ -563,46 +592,45 @@ def petsc_dae_by_time_element(
         initial_solver_obj = pyo.SolverFactory(
             initial_solver, options=initial_solver_options
         )
-        # list of constraints to add to the initial condition problem
-        if initial_constraints is None:
-            initial_constraints = []
-
-        if detect_initial:
-            # If detect_initial, solve the non-time-indexed variables and
-            # constraints with the initial conditions
-            rcset = ComponentSet(regular_cons)
-            icset = ComponentSet(initial_constraints)
-            initial_constraints = list(icset | rcset)
-
-        with TemporarySubsystemManager(to_deactivate=tdisc):
-            constraints = [
-                con[t0] for con in time_cons if t0 in con
-            ] + initial_constraints
-            variables = [var[t0] for var in time_vars] + initial_variables
-            if len(constraints) > 0:
-                # if the initial condition is specified and there are no
-                # initial constraints, don't try to solve.
-                t_block = create_subsystem_block(
-                    constraints,
-                    variables,
+        try:
+            init_subsystem = get_initial_condition_problem(
+                model=m,
+                time=time,
+                t_initial=t0,
+                initial_constraints=initial_constraints,
+                initial_variables=initial_variables,
+                detect_initial=detect_initial,
+                flattened_problem=flattened_problem,
+            )
+        except RuntimeError as err:
+            if "Zero constraints" in err.message:
+                init_subsystem=None
+            else:
+                raise err
+        if init_subsystem is not None and number_activated_constraints(init_subsystem) > 0:
+            dof = degrees_of_freedom(init_subsystem)
+            if dof > 0:
+                raise RuntimeError(
+                    "Expected zero degrees of freedom in initial condition problem, but "
+                    "instead found {dof} degrees of freedom. You can use the function "
+                    "get_initial_condition_problem to get a copy of the initial "
+                    "condition problem to debug yourself."
                 )
-                # set up the scaling factor suffix
-                _sub_problem_scaling_suffix(m, t_block)
-                with idaeslog.solver_log(solve_log, idaeslog.INFO) as slc:
-                    res = initial_solver_obj.solve(
-                        t_block,
-                        tee=slc.tee,
-                        symbolic_solver_labels=symbolic_solver_labels,
-                    )
-                res_list.append(res)
+            with idaeslog.solver_log(solve_log, idaeslog.INFO) as slc:
+                res = initial_solver_obj.solve(
+                    init_subsystem,
+                    tee=slc.tee,
+                    symbolic_solver_labels=symbolic_solver_labels,
+                )
+            res_list.append(res)
 
     tprev = t0
     tj = previous_trajectory
     if tj is not None:
-        variables_prev = [var[t0] for var in time_vars]
+        variables_prev = [var[t0] for var in flattened_problem["time_variables"]]
 
     with TemporarySubsystemManager(
-        to_deactivate=tdisc,
+        to_deactivate=flattened_problem["discretization_equations"],
         to_fix=initial_variables,
     ):
         # Solver time steps
@@ -611,8 +639,8 @@ def petsc_dae_by_time_element(
             if t == between.first():
                 # t == between.first() was handled above
                 continue
-            constraints = [con[t] for con in time_cons if t in con]
-            variables = [var[t] for var in time_vars]
+            constraints = [con[t] for con in flattened_problem["time_constraints"] if t in con]
+            variables = [var[t] for var in flattened_problem["time_variables"]]
             # Create a temporary block with references to original constraints
             # and variables so we can integrate this "subsystem" without
             # altering the rest of the model.
@@ -635,7 +663,7 @@ def petsc_dae_by_time_element(
             # Set up the scaling factor suffix
             _sub_problem_scaling_suffix(m, t_block)
             # Take initial conditions for this step from the result of previous
-            _copy_time(time_vars, tprev, t)
+            _copy_time(flattened_problem["time_variables"], tprev, t)
             with idaeslog.solver_log(solve_log, idaeslog.INFO) as slc:
                 res = solver_dae.solve(
                     t_block,
@@ -696,7 +724,7 @@ def petsc_dae_by_time_element(
                 for t in itime:
                     timevar[t] = t
                 no_repeat.add(id(timevar[tlast]))
-            for var in time_vars:
+            for var in flattened_problem["time_variables"]:
                 if id(var[tlast]) in no_repeat:
                     continue
                 no_repeat.add(id(var[tlast]))
@@ -772,7 +800,12 @@ def calculate_time_derivatives(m, time, between=None):
                             if t == time.first() or t == time.last():
                                 pass
                             else:
-                                raise err
+                                raise RuntimeError(
+                                    f"Discretization equation corresponding to {deriv[t].name} does not "
+                                    f"exist at time {t}, so derivative values cannot be calculated. Check "
+                                    "to see if it if it is supposed to exist. For certain collocation methods "
+                                    "(for example, Lagrange-Legendre) it does not."
+                                ) from err
                         except ValueError as err:
                             # At edges of between, it's unclear which adjacent
                             # values of state variables have been populated.
@@ -1023,3 +1056,83 @@ def from_json(self, pth):
             with open(pth, "r") as fp:
                 self.vecs = json.load(fp)
         self.time = self.vecs["_time"]
+
+def get_flattened_problem(model, time, representative_time):
+    regular_vars, time_vars = flatten_dae_components(
+        model, time, pyo.Var, active=True, indices=(representative_time,)
+    )
+    regular_cons, time_cons = flatten_dae_components(
+        model, time, pyo.Constraint, active=True, indices=(representative_time,)
+    )
+    tdisc = ComponentSet(get_discretization_constraint_data(model, time))
+    flattened_problem = {
+        "timeless_variables": regular_vars,
+        "time_variables": time_vars,
+        "timeless_constraints": regular_cons,
+        "time_constraints": time_cons,
+        "discretization_equations": tdisc,
+    }
+    return flattened_problem
+
+
+def get_initial_condition_problem(
+        model,
+        time,
+        t_initial,
+        representative_time=None,
+        initial_constraints=None,
+        initial_variables=None,
+        detect_initial=True,
+        flattened_problem=None,
+    ):
+        if initial_variables is None:
+            initial_variables = []
+        # list of constraints to add to the initial condition problem
+        if initial_constraints is None:
+            initial_constraints = []
+
+        if flattened_problem is None:
+            if representative_time is None:
+                raise RuntimeError("The user must supply either the flattened problem or a representative time.")
+            flattened_problem = get_flattened_problem(
+                model=model,
+                time=time,
+                representative_time=representative_time
+            )
+
+        if detect_initial:
+            rvset = ComponentSet(flattened_problem["timeless_variables"])
+            ivset = ComponentSet(initial_variables)
+            initial_variables = list(ivset | rvset)
+            # If detect_initial, solve the non-time-indexed variables and
+            # constraints with the initial conditions
+            const_no_t_set = ComponentSet(flattened_problem["timeless_constraints"])
+            const_init_set = ComponentSet(initial_constraints)
+            initial_constraints = list(const_no_t_set | const_init_set)
+
+
+
+        constraints = [
+            con[t_initial] 
+            for con in flattened_problem["time_constraints"]
+            if t_initial in con
+            and con[t_initial] not in flattened_problem["discretization_equations"]
+        ] + initial_constraints
+        variables = [
+            var[t_initial] 
+            for var in flattened_problem["time_variables"]
+        ] + initial_variables
+
+        if len(constraints) <= 0:
+            raise RuntimeError(
+                "Zero constraints in initial condition problem, therefore "
+                "there is no block to return."
+            )
+        # if the initial condition is specified and there are no
+        # initial constraints, don't try to solve.
+        subsystem = create_subsystem_block(
+            constraints,
+            variables,
+        )
+        _sub_problem_scaling_suffix(model, subsystem)
+        return subsystem
\ No newline at end of file
diff --git a/idaes/core/solvers/tests/test_petsc.py b/idaes/core/solvers/tests/test_petsc.py
index f19f51acb8..3bc1459b23 100644
--- a/idaes/core/solvers/tests/test_petsc.py
+++ b/idaes/core/solvers/tests/test_petsc.py
@@ -21,6 +21,7 @@
 import pyomo.dae as pyodae
 from pyomo.util.subsystems import create_subsystem_block
 from idaes.core.solvers import petsc
+import idaes.core.util.model_statistics as mstat
 import idaes.logger as idaeslog
 
 
@@ -835,10 +836,11 @@ def test_petsc_traj_previous():
     tj0 = res0.trajectory
 
     # Next do it in two steps
+    t4 = m.t.at(4)
     res = petsc.petsc_dae_by_time_element(
         m,
         time=m.t,
-        between=[m.t.first(), m.t.at(4)],
+        between=[m.t.first(), t4],
         ts_options={
             "--ts_type": "cn",  # Crank–Nicolson
             "--ts_adapt_type": "basic",
@@ -848,11 +850,16 @@ def test_petsc_traj_previous():
     )
     # Fix initial condition for second leg of trajectory
     for j in range(1, 6):
-        m.y[m.t.at(4), j].fix()
+        m.y[t4, j].fix()
+    m.eq_ydot1[t4].deactivate()
+    m.eq_ydot2[t4].deactivate()
+    m.eq_ydot3[t4].deactivate()
+    m.eq_ydot4[t4].deactivate()
+    m.eq_ydot5[t4].deactivate()
     res = petsc.petsc_dae_by_time_element(
         m,
         time=m.t,
-        between=[m.t.at(4), m.t.last()],
+        between=[t4, m.t.last()],
         ts_options={
             "--ts_type": "cn",  # Crank–Nicolson
             "--ts_adapt_type": "basic",
@@ -860,6 +867,17 @@ def test_petsc_traj_previous():
             "--ts_save_trajectory": 1,
         },
         previous_trajectory=res.trajectory,
+        initial_solver="ipopt",
+        initial_solver_options={
+            # 'bound_push' : 1e-22,
+            'linear_solver': 'ma57',
+            'OF_ma57_automatic_scaling': 'yes',
+            'max_iter': 300,
+            'tol': 1e-6,
+            'halt_on_ampl_error': 'yes',
+            # 'mu_strategy': 'monotone',
+        },
+        # skip_initial=True,
         # interpolate=False
     )
 
@@ -1232,40 +1250,214 @@ def diff_eq_rule(m, t):
 
 
 # This test fails because the PETSc interface doesn't work with Lagrange-Legendre collocation
-# @pytest.mark.unit
-# @pytest.mark.skipif(not petsc.petsc_available(), reason="PETSc solver not available")
-# def test_calculate_derivatives_collocation_ll():
-#     m = pyo.ConcreteModel()
-
-#     m.time = pyodae.ContinuousSet(initialize=(0.0, 1.0, 2.0))
-#     m.x = pyo.Var(m.time)
-#     m.u = pyo.Var(m.time)
-#     m.dxdt = pyodae.DerivativeVar(m.x, wrt=m.time)
-
-#     def diff_eq_rule(m, t):
-#         return m.dxdt[t] == m.x[t] ** 2 - m.u[t]
-
-#     m.diff_eq = pyo.Constraint(m.time, rule=diff_eq_rule)
-
-#     discretizer = pyo.TransformationFactory("dae.collocation")
-#     discretizer.apply_to(m, nfe=2, scheme="LAGRANGE-LEGENDRE")
-
-#     m.u.fix(1.0)
-#     m.x[0].fix(0.0)
-
-#     res = petsc.petsc_dae_by_time_element(
-#         m,
-#         time=m.time,
-#         between=[0.0, 2.0],
-#         ts_options={
-#             "--ts_type": "beuler",
-#             "--ts_dt": 3e-2,
-#         },
-#         skip_initial=True,  # With u and x fixed, no variables to solve for at t0
-#         calculate_derivatives=True,
-#     )
-
-#     assert m.dxdt[0].value is None
-#     # It should backfill values for interpolated points like t=1
-#     assert pyo.value(m.dxdt[1]) == pytest.approx(-0.3787483909827891, rel=1e-3)
-#     assert pyo.value(m.dxdt[2]) == pytest.approx(-0.07952012649572815, rel=1e-3)
+@pytest.mark.unit
+@pytest.mark.skipif(not petsc.petsc_available(), reason="PETSc solver not available")
+def test_calculate_derivatives_collocation_ll():
+    m = pyo.ConcreteModel()
+
+    m.time = pyodae.ContinuousSet(initialize=(0.0, 1.0, 2.0))
+    m.x = pyo.Var(m.time)
+    m.u = pyo.Var(m.time)
+    m.dxdt = pyodae.DerivativeVar(m.x, wrt=m.time)
+
+    def diff_eq_rule(m, t):
+        return m.dxdt[t] == m.x[t] ** 2 - m.u[t]
+
+    m.diff_eq = pyo.Constraint(m.time, rule=diff_eq_rule)
+
+    discretizer = pyo.TransformationFactory("dae.collocation")
+    discretizer.apply_to(m, nfe=2, scheme="LAGRANGE-LEGENDRE")
+
+    m.u.fix(1.0)
+    m.x[0].fix(0.0)
+    with pytest.raises(
+        RuntimeError,
+        match=re.escape(
+            "Discretization equation corresponding to dxdt[1.0] does not exist at time 1.0, "
+            "so derivative values cannot be calculated. Check to see if it if it is supposed "
+            "to exist. For certain collocation methods (for example, Lagrange-Legendre) "
+            "it does not."
+        ),
+    ):
+        res = petsc.petsc_dae_by_time_element(
+            m,
+            time=m.time,
+            between=[0.0, 2.0],
+            ts_options={
+                "--ts_type": "beuler",
+                "--ts_dt": 3e-2,
+            },
+            skip_initial=True,  # With u and x fixed, no variables to solve for at t0
+            calculate_derivatives=True,
+        )
+
+@pytest.mark.unit
+def test_get_initial_condition_problem_no_active_constraints():
+    m = pyo.ConcreteModel()
+
+    m.time = pyodae.ContinuousSet(initialize=(0.0, 1.0, 2.0))
+    m.x = pyo.Var(m.time)
+    m.u = pyo.Var(m.time)
+    m.dxdt = pyodae.DerivativeVar(m.x, wrt=m.time)
+
+    def diff_eq_rule(m, t):
+        return m.dxdt[t] == m.x[t] ** 2 - m.u[t]
+
+    m.diff_eq = pyo.Constraint(m.time, rule=diff_eq_rule)
+    m.diff_eq[0].deactivate()
+
+    m.u.fix(1.0)
+    m.x[0].fix(0.0)
+
+    discretizer = pyo.TransformationFactory("dae.finite_difference")
+    discretizer.apply_to(m, nfe=2, scheme="BACKWARD")
+
+    init_subsystem = petsc.get_initial_condition_problem(
+        model=m,
+        time=m.time,
+        t_initial=m.time.at(1),
+        representative_time=m.time.at(2)
+    )
+    assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
+    local_var_name_list = [var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)]
+    assert set(local_var_name_list) == {'x[0.0]', 'u[0.0]', 'dxdt[0.0]'}
+
+    assert mstat.number_activated_constraints(init_subsystem) == 0
+    local_con_name_list = [con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)]
+    assert local_con_name_list[0] == "diff_eq[0.0]"
+
+@pytest.mark.unit
+def test_get_initial_condition_problem_no_active_constraints_not_t0():
+    m = pyo.ConcreteModel()
+
+    m.time = pyodae.ContinuousSet(initialize=(0.0, 1.0, 2.0))
+    m.x = pyo.Var(m.time)
+    m.u = pyo.Var(m.time)
+    m.dxdt = pyodae.DerivativeVar(m.x, wrt=m.time)
+
+    def diff_eq_rule(m, t):
+        return m.dxdt[t] == m.x[t] ** 2 - m.u[t]
+
+    m.diff_eq = pyo.Constraint(m.time, rule=diff_eq_rule)
+    m.diff_eq[1].deactivate()
+
+    m.u.fix(1.0)
+    m.x[1].fix(0.0)
+
+    discretizer = pyo.TransformationFactory("dae.finite_difference")
+    discretizer.apply_to(m, nfe=2, scheme="BACKWARD")
+
+    init_subsystem = petsc.get_initial_condition_problem(
+        model=m,
+        time=m.time,
+        t_initial=m.time.at(2),
+        representative_time=m.time.at(3)
+    )
+    assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
+    local_var_name_list = [var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)]
+    assert set(local_var_name_list) == {'x[1.0]', 'u[1.0]', 'dxdt[1.0]'}
+
+    assert mstat.number_activated_constraints(init_subsystem) == 0
+    local_con_name_list = [con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)]
+    assert local_con_name_list[0] == "diff_eq[1.0]"
+
+@pytest.mark.unit
+def test_get_initial_condition_problem_at_time_no_active_constraints():
+    m = pyo.ConcreteModel()
+
+    m.time = pyodae.ContinuousSet(initialize=(0.0, 1.0, 2.0))
+    m.x = pyo.Var(m.time)
+    m.u = pyo.Var(m.time)
+    m.dxdt = pyodae.DerivativeVar(m.x, wrt=m.time)
+
+    def diff_eq_rule(m, t):
+        return m.dxdt[t] == m.x[t] ** 2 - m.u[t]
+
+    m.diff_eq = pyo.Constraint(m.time, rule=diff_eq_rule)
+    m.diff_eq[0].deactivate()
+
+    discretizer = pyo.TransformationFactory("dae.finite_difference")
+    discretizer.apply_to(m, nfe=2, scheme="BACKWARD")
+
+    init_subsystem = petsc.get_initial_condition_problem(
+        model=m,
+        time=m.time,
+        t_initial=m.time.at(1),
+        representative_time=m.time.at(2)
+    )
+    assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
+    local_var_name_list = [var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)]
+    assert set(local_var_name_list) == {'x[0.0]', 'u[0.0]', 'dxdt[0.0]'}
+
+    assert mstat.number_activated_constraints(init_subsystem) == 0
+    local_con_name_list = [con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)]
+    assert local_con_name_list[0] == "diff_eq[0.0]"
+
+@pytest.mark.unit
+def test_get_initial_condition_problem_non_time_indexed_constraint():
+    m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(nfe=10)
+    init_subsystem = petsc.get_initial_condition_problem(
+        model=m, time=m.t, t_initial=m.t.at(1), representative_time=m.t.at(2)
+    )
+    
+    unfixed_var_name_set = set(var.name for var in mstat.unfixed_variables_set(init_subsystem))
+    assert unfixed_var_name_set == {'y[0,6]', 'ydot[0,1]', 'ydot[0,2]', 'ydot[0,3]', 'ydot[0,4]', 'ydot[0,5]', 'r[0,1]', 'r[0,2]', 'r[0,3]', 'r[0,4]', 'r[0,5]', 'Fin[0]', 'H'}
+
+    total_var_name_set = set(var.name for var in mstat.variables_set(init_subsystem))
+    assert total_var_name_set == {'r[0,1]', 'r[0,2]', 'ydot[0,1]', 'ydot[0,4]', 'H', 'y[0,2]', 'ydot[0,3]', 'y[0,3]', 'y[0,5]', 'r[0,3]', 'r[0,5]', 'Fin[0]', 'y[0,1]', 'ydot[0,2]', 'ydot[0,6]', 'r[0,4]', 'y[0,4]', 'y[0,6]', 'ydot[0,5]'}
+
+    active_con_name_set = set(con.name for con in mstat.activated_constraints_set(init_subsystem))
+    assert active_con_name_set == {'eq_y6[0]', 'eq_r1[0]', 'eq_r2[0]', 'eq_r3[0]', 'eq_r4[0]', 'eq_r5[0]', 'eq_Fin[0]', 'H_eqn'}
+
+    total_con_name_set = set(con.name for con in mstat.total_constraints_set(init_subsystem))
+    assert total_con_name_set == {'eq_r5[0]', 'eq_y6[0]', 'eq_r4[0]', 'eq_Fin[0]', 'H_eqn', 'eq_ydot5[0]', 'eq_ydot1[0]', 'eq_r1[0]', 'eq_ydot4[0]', 'eq_r3[0]', 'eq_r2[0]', 'eq_ydot2[0]', 'eq_ydot3[0]'}
+
+@pytest.mark.unit
+def test_get_initial_condition_problem_non_time_indexed_constraint_not_t0():
+    m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(nfe=10)
+    init_subsystem = petsc.get_initial_condition_problem(
+        model=m, time=m.t, t_initial=m.t.at(4), representative_time=m.t.at(5)
+    )
+    t4 = m.t.at(4)
+    for j in range(1, 6):
+        m.y[t4, j].fix()
+    m.eq_ydot1[t4].deactivate()
+    m.eq_ydot2[t4].deactivate()
+    m.eq_ydot3[t4].deactivate()
+    m.eq_ydot4[t4].deactivate()
+    m.eq_ydot5[t4].deactivate()
+
+    unfixed_var_name_set = set(var.name for var in mstat.unfixed_variables_set(init_subsystem))
+    assert unfixed_var_name_set == {'y[54.0,6]', 'ydot[54.0,1]', 'ydot[54.0,2]', 'ydot[54.0,3]', 'ydot[54.0,4]', 'ydot[54.0,5]', 'r[54.0,1]', 'r[54.0,2]', 'r[54.0,3]', 'r[54.0,4]', 'r[54.0,5]', 'Fin[54.0]', 'H'}
+
+    total_var_name_set = set(var.name for var in mstat.variables_set(init_subsystem))
+    assert total_var_name_set == {'r[54.0,1]', 'r[54.0,2]', 'ydot[54.0,1]', 'ydot[54.0,4]', 'H', 'y[54.0,2]', 'ydot[54.0,3]', 'y[54.0,3]', 'y[54.0,5]', 'r[54.0,3]', 'r[54.0,5]', 'Fin[54.0]', 'y[54.0,1]', 'ydot[54.0,2]', 'ydot[54.0,6]', 'r[54.0,4]', 'y[54.0,4]', 'y[54.0,6]', 'ydot[54.0,5]'}
+
+    active_con_name_set = set(con.name for con in mstat.activated_constraints_set(init_subsystem))
+    assert active_con_name_set == {'eq_y6[54.0]', 'eq_r1[54.0]', 'eq_r2[54.0]', 'eq_r3[54.0]', 'eq_r4[54.0]', 'eq_r5[54.0]', 'eq_Fin[54.0]', 'H_eqn'}
+
+    total_con_name_set = set(con.name for con in mstat.total_constraints_set(init_subsystem))
+    assert total_con_name_set == {'eq_r5[54.0]', 'eq_y6[54.0]', 'eq_r4[54.0]', 'eq_Fin[54.0]', 'H_eqn', 'eq_ydot5[54.0]', 'eq_ydot1[54.0]', 'eq_r1[54.0]', 'eq_ydot4[54.0]', 'eq_r3[54.0]', 'eq_r2[54.0]', 'eq_ydot2[54.0]', 'eq_ydot3[54.0]'}
+
+
+if __name__ == "__main__":
+    # Will clean this up after review
+    m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(nfe=10)
+
+    # init_subsystem = petsc.get_initial_condition_problem_at_time(
+    #     model=m, time=m.t, t_initial=m.t.at(1), representative_time=m.t.at(2)
+    # )
+
+    t4 = m.t.at(4)
+    for j in range(1, 6):
+        m.y[t4, j].fix()
+    m.eq_ydot1[t4].deactivate()
+    m.eq_ydot2[t4].deactivate()
+    m.eq_ydot3[t4].deactivate()
+    m.eq_ydot4[t4].deactivate()
+    m.eq_ydot5[t4].deactivate()
+    init_subsystem = petsc.get_initial_condition_problem(model=m, time=m.t, t_initial=m.t.at(4), representative_time=m.t.at(5))
+    from idaes.core.util.model_diagnostics import DiagnosticsToolbox
+    dt=DiagnosticsToolbox(init_subsystem)
+    dt.report_structural_issues()
+    print("ok, boomer")

From 850605a1cfd62e1d124f0b5d2950f8c8e5fc3a92 Mon Sep 17 00:00:00 2001
From: Doug A <douglas.allan@keylogic.com>
Date: Wed, 26 Jun 2024 13:07:22 -0400
Subject: [PATCH 2/4] Get rid of pytest warnings about returning from tests

---
 idaes/models/control/tests/test_steam_tank_pressure.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/idaes/models/control/tests/test_steam_tank_pressure.py b/idaes/models/control/tests/test_steam_tank_pressure.py
index 76ee0f1223..09521913cd 100644
--- a/idaes/models/control/tests/test_steam_tank_pressure.py
+++ b/idaes/models/control/tests/test_steam_tank_pressure.py
@@ -293,8 +293,6 @@ def test_inlet_disturbance():
         m_dynamic.fs.valve_1.valve_opening[m_dynamic.fs.time.last()]
     ) == pytest.approx(s2_valve, abs=0.001)
 
-    return m_dynamic, solver
-
 
 @pytest.mark.skipif(not helmholtz_available(), reason="General Helmholtz not available")
 @pytest.mark.integration
@@ -336,8 +334,6 @@ def test_inlet_disturbance_derivative_on_error():
         m_dynamic.fs.valve_1.valve_opening[m_dynamic.fs.time.last()]
     ) == pytest.approx(s2_valve, abs=0.001)
 
-    return m_dynamic, solver
-
 
 @pytest.mark.skipif(not helmholtz_available(), reason="General Helmholtz not available")
 @pytest.mark.integration
@@ -384,8 +380,6 @@ def test_setpoint_change_derivative_on_pv():
         m_dynamic.fs.valve_1.valve_opening[m_dynamic.fs.time.at(31)]
     ) == pytest.approx(0.78949, abs=0.001)
 
-    return m_dynamic, solver
-
 
 @pytest.mark.skipif(not helmholtz_available(), reason="General Helmholtz not available")
 @pytest.mark.integration
@@ -432,8 +426,6 @@ def test_setpoint_change_derivative_on_error():
         m_dynamic.fs.valve_1.valve_opening[m_dynamic.fs.time.at(31)]
     ) == pytest.approx(0.93196, abs=0.001)
 
-    return m_dynamic, solver
-
 
 if __name__ == "__main__":
     m, solver = test_setpoint_change_derivative_on_pv()

From 01f058b0b80a0fbc91138798d0aa1b313c0d765c Mon Sep 17 00:00:00 2001
From: Doug A <douglas.allan@keylogic.com>
Date: Wed, 26 Jun 2024 15:15:12 -0400
Subject: [PATCH 3/4] before black

---
 docs/reference_guides/core/solvers.rst |  12 ++
 idaes/core/solvers/petsc.py            | 152 ++++++++++++++++---------
 idaes/core/solvers/tests/test_petsc.py |  75 ++++++++++--
 3 files changed, 179 insertions(+), 60 deletions(-)

diff --git a/docs/reference_guides/core/solvers.rst b/docs/reference_guides/core/solvers.rst
index b0479bef5c..bdf6eea049 100644
--- a/docs/reference_guides/core/solvers.rst
+++ b/docs/reference_guides/core/solvers.rst
@@ -186,6 +186,18 @@ to load, save, and interpolate.
 .. autoclass:: idaes.core.solvers.petsc.PetscTrajectory
     :members:
 
+Diagnostics
+"""""""""""
+
+Often, the first formulation of a dynamic problem is not successful, suffering form
+structural or numerical singularity or other modeling issues. IDAES features a number
+of :ref:`model diagnostics functions<explanations/model_diagnostics/index:Model Diagnostics Workflow>`
+that can be used to help troubleshoot issues. In order to apply these functions to the initial
+condition subproblem that is solved as part of the PETSc interface, the following 
+function returns a Pyomo block for that problem:
+
+.. autofunction:: idaes.core.solvers.petsc.get_initial_condition_problem
+
 Using TS Solvers without Pyomo.DAE
 """"""""""""""""""""""""""""""""""
 
diff --git a/idaes/core/solvers/petsc.py b/idaes/core/solvers/petsc.py
index fdce3adf60..7129ccbee5 100644
--- a/idaes/core/solvers/petsc.py
+++ b/idaes/core/solvers/petsc.py
@@ -567,7 +567,7 @@ def petsc_dae_by_time_element(
     elif representative_time not in between:
         raise RuntimeError("representative_time is not element of between.")
 
-    flattened_problem  = get_flattened_problem(
+    flattened_problem  = _get_flattened_problem(
         model=m,
         time=time,
         representative_time=representative_time
@@ -596,7 +596,7 @@ def petsc_dae_by_time_element(
             init_subsystem = get_initial_condition_problem(
                 model=m,
                 time=time,
-                t_initial=t0,
+                initial_time=t0,
                 initial_constraints=initial_constraints,
                 initial_variables=initial_variables,
                 detect_initial=detect_initial,
@@ -1057,7 +1057,23 @@ def from_json(self, pth):
                 self.vecs = json.load(fp)
         self.time = self.vecs["_time"]
 
-def get_flattened_problem(model, time, representative_time):
+def _get_flattened_problem(model, time, representative_time):
+    """
+    Helper function for petsc_dae_by_time_element and get_initial_condition_problem.
+    Gets a view of the problem flattened so time is the only explicit index
+
+    Args:
+        model (Block): Pyomo model to solve
+        time (ContinuousSet): Time set
+        representative_time (Element of time): A timepoint at which the DAE system is at its "normal" state
+            after the constraints and variables associated with the initial time point
+            have passed. 
+
+    Returns (dictionary):
+        Dictionary containing lists of variables and constraints indexed by times and lists
+        of those unindexed by time. Also a list of discretization equations so they can be 
+        deactivated in the problem passed to PETSc.
+    """
     regular_vars, time_vars = flatten_dae_components(
         model, time, pyo.Var, active=True, indices=(representative_time,)
     )
@@ -1078,61 +1094,91 @@ def get_flattened_problem(model, time, representative_time):
 def get_initial_condition_problem(
         model,
         time,
-        t_initial,
+        initial_time,
         representative_time=None,
         initial_constraints=None,
         initial_variables=None,
         detect_initial=True,
         flattened_problem=None,
     ):
-        if initial_variables is None:
-            initial_variables = []
-        # list of constraints to add to the initial condition problem
-        if initial_constraints is None:
-            initial_constraints = []
-
-        if flattened_problem is None:
-            if representative_time is None:
-                raise RuntimeError("The user must supply either the flattened problem or a representative time.")
-            flattened_problem = get_flattened_problem(
-                model=model,
-                time=time,
-                representative_time=representative_time
-            )
+    """
+    Solve a DAE problem step by step using the PETSc DAE solver.  This
+    integrates from one time point to the next.
 
-        if detect_initial:
-            rvset = ComponentSet(flattened_problem["timeless_variables"])
-            ivset = ComponentSet(initial_variables)
-            initial_variables = list(ivset | rvset)
-            # If detect_initial, solve the non-time-indexed variables and
-            # constraints with the initial conditions
-            const_no_t_set = ComponentSet(flattened_problem["timeless_constraints"])
-            const_init_set = ComponentSet(initial_constraints)
-            initial_constraints = list(const_no_t_set | const_init_set)
-
-
-
-        constraints = [
-            con[t_initial] 
-            for con in flattened_problem["time_constraints"]
-            if t_initial in con
-            and con[t_initial] not in flattened_problem["discretization_equations"]
-        ] + initial_constraints
-        variables = [
-            var[t_initial] 
-            for var in flattened_problem["time_variables"]
-        ] + initial_variables
-
-        if len(constraints) <= 0:
-            raise RuntimeError(
-                "Zero constraints in initial condition problem, therefore "
-                "there is no block to return."
-            )
-        # if the initial condition is specified and there are no
-        # initial constraints, don't try to solve.
-        subsystem = create_subsystem_block(
-            constraints,
-            variables,
+    Args:
+        model (Block): Pyomo model to solve
+        time (ContinuousSet): Time set
+        initial_time (Element of time): The timepoint to return initial condition problem for
+        representative_time (Element of time): A timepoint at which the DAE system is at its "normal" state
+            after the constraints and variables associated with the initial time point
+            have passed. Typically the next element of time after initial_time. Not needed
+            if flattened_problem is provided.
+        initial_constraints (list): Constraints to solve in the initial
+            condition solve step.  Since the time-indexed constraints are picked
+            up automatically, this generally includes non-time-indexed
+            constraints.
+        initial_variables (list): This is a list of variables to fix after the
+            initial condition solve step.  If these variables were originally
+            unfixed, they will be unfixed at the end of the solve. This usually
+            includes non-time-indexed variables that are calculated along with
+            the initial conditions.
+        detect_initial (bool): If True, add non-time-indexed variables and
+            constraints to initial_variables and initial_constraints.
+        flattened_problem (dict): Dictionary returned by get_flattened_problem.
+            If not provided, get_flattened_problem will be called at representative_time.
+
+    Returns (Pyomo Block):
+        Block containing References to variables and constraints used in initial condition
+        problem, ready to be solved as the initial condition problem.
+    """
+    if initial_variables is None:
+        initial_variables = []
+    # list of constraints to add to the initial condition problem
+    if initial_constraints is None:
+        initial_constraints = []
+
+    if flattened_problem is None:
+        if representative_time is None:
+            raise RuntimeError("The user must supply either the flattened problem or a representative time.")
+        flattened_problem = _get_flattened_problem(
+            model=model,
+            time=time,
+            representative_time=representative_time
         )
-        _sub_problem_scaling_suffix(model, subsystem)
-        return subsystem
\ No newline at end of file
+
+    if detect_initial:
+        rvset = ComponentSet(flattened_problem["timeless_variables"])
+        ivset = ComponentSet(initial_variables)
+        initial_variables = list(ivset | rvset)
+        # If detect_initial, solve the non-time-indexed variables and
+        # constraints with the initial conditions
+        const_no_t_set = ComponentSet(flattened_problem["timeless_constraints"])
+        const_init_set = ComponentSet(initial_constraints)
+        initial_constraints = list(const_no_t_set | const_init_set)
+
+
+
+    constraints = [
+        con[initial_time] 
+        for con in flattened_problem["time_constraints"]
+        if initial_time in con
+        and con[initial_time] not in flattened_problem["discretization_equations"]
+    ] + initial_constraints
+    variables = [
+        var[initial_time] 
+        for var in flattened_problem["time_variables"]
+    ] + initial_variables
+
+    if len(constraints) <= 0:
+        raise RuntimeError(
+            "Zero constraints in initial condition problem, therefore "
+            "there is no block to return."
+        )
+    # if the initial condition is specified and there are no
+    # initial constraints, don't try to solve.
+    subsystem = create_subsystem_block(
+        constraints,
+        variables,
+    )
+    _sub_problem_scaling_suffix(model, subsystem)
+    return subsystem
\ No newline at end of file
diff --git a/idaes/core/solvers/tests/test_petsc.py b/idaes/core/solvers/tests/test_petsc.py
index 3bc1459b23..846ef067c4 100644
--- a/idaes/core/solvers/tests/test_petsc.py
+++ b/idaes/core/solvers/tests/test_petsc.py
@@ -812,6 +812,68 @@ def test_petsc_skip_initial_solve():
     assert pytest.approx(y5, rel=1e-3) == pyo.value(m.y[m.t.last(), 5])
     assert pytest.approx(y6, rel=1e-3) == pyo.value(m.y[m.t.last(), 6])
 
+@pytest.mark.unit
+@pytest.mark.skipif(not petsc.petsc_available(), reason="PETSc solver not available")
+def test_petsc_collocation_lagrange_radau():
+    """
+    Ensure that PETSc works with Lagrange-Radau collocation
+    """
+    m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(
+        transformation_method="dae.collocation",
+        scheme="LAGRANGE-RADAU",
+        ncp=3
+    )
+
+    res = petsc.petsc_dae_by_time_element(
+        m,
+        time=m.t,
+        ts_options={
+            "--ts_type": "cn",  # Crank–Nicolson
+            "--ts_adapt_type": "basic",
+            "--ts_dt": 0.01,
+            "--ts_save_trajectory": 1,
+            "--ts_trajectory_type": "visualization",
+        },
+    )
+
+    assert pytest.approx(y1, rel=1e-3) == pyo.value(m.y[m.t.last(), 1])
+    assert pytest.approx(y2, rel=1e-3) == pyo.value(m.y[m.t.last(), 2])
+    assert pytest.approx(y3, rel=1e-3) == pyo.value(m.y[m.t.last(), 3])
+    assert pytest.approx(y4, rel=1e-3) == pyo.value(m.y[m.t.last(), 4])
+    assert pytest.approx(y5, rel=1e-3) == pyo.value(m.y[m.t.last(), 5])
+    assert pytest.approx(y6, rel=1e-3) == pyo.value(m.y[m.t.last(), 6])
+
+@pytest.mark.unit
+@pytest.mark.skipif(not petsc.petsc_available(), reason="PETSc solver not available")
+def test_petsc_collocation_lagrange_legendre():
+    """
+    Ensure that PETSc works with Lagrange-Legendre collocation
+    """
+    m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(
+        transformation_method="dae.collocation",
+        scheme="LAGRANGE-LEGENDRE",
+        ncp=3
+    )
+
+    res = petsc.petsc_dae_by_time_element(
+        m,
+        time=m.t,
+        ts_options={
+            "--ts_type": "cn",  # Crank–Nicolson
+            "--ts_adapt_type": "basic",
+            "--ts_dt": 0.01,
+            "--ts_save_trajectory": 1,
+            "--ts_trajectory_type": "visualization",
+        },
+    )
+
+    assert pytest.approx(y1, rel=1e-3) == pyo.value(m.y[m.t.last(), 1])
+    assert pytest.approx(y2, rel=1e-3) == pyo.value(m.y[m.t.last(), 2])
+    assert pytest.approx(y3, rel=1e-3) == pyo.value(m.y[m.t.last(), 3])
+    assert pytest.approx(y4, rel=1e-3) == pyo.value(m.y[m.t.last(), 4])
+    assert pytest.approx(y5, rel=1e-3) == pyo.value(m.y[m.t.last(), 5])
+    assert pytest.approx(y6, rel=1e-3) == pyo.value(m.y[m.t.last(), 6])
+
 
 @pytest.mark.unit
 @pytest.mark.skipif(not petsc.petsc_available(), reason="PETSc solver not available")
@@ -1315,7 +1377,7 @@ def diff_eq_rule(m, t):
     init_subsystem = petsc.get_initial_condition_problem(
         model=m,
         time=m.time,
-        t_initial=m.time.at(1),
+        initial_time=m.time.at(1),
         representative_time=m.time.at(2)
     )
     assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
@@ -1350,7 +1412,7 @@ def diff_eq_rule(m, t):
     init_subsystem = petsc.get_initial_condition_problem(
         model=m,
         time=m.time,
-        t_initial=m.time.at(2),
+        initial_time=m.time.at(2),
         representative_time=m.time.at(3)
     )
     assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
@@ -1382,7 +1444,7 @@ def diff_eq_rule(m, t):
     init_subsystem = petsc.get_initial_condition_problem(
         model=m,
         time=m.time,
-        t_initial=m.time.at(1),
+        initial_time=m.time.at(1),
         representative_time=m.time.at(2)
     )
     assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
@@ -1397,7 +1459,7 @@ def diff_eq_rule(m, t):
 def test_get_initial_condition_problem_non_time_indexed_constraint():
     m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(nfe=10)
     init_subsystem = petsc.get_initial_condition_problem(
-        model=m, time=m.t, t_initial=m.t.at(1), representative_time=m.t.at(2)
+        model=m, time=m.t, initial_time=m.t.at(1), representative_time=m.t.at(2)
     )
     
     unfixed_var_name_set = set(var.name for var in mstat.unfixed_variables_set(init_subsystem))
@@ -1416,7 +1478,7 @@ def test_get_initial_condition_problem_non_time_indexed_constraint():
 def test_get_initial_condition_problem_non_time_indexed_constraint_not_t0():
     m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(nfe=10)
     init_subsystem = petsc.get_initial_condition_problem(
-        model=m, time=m.t, t_initial=m.t.at(4), representative_time=m.t.at(5)
+        model=m, time=m.t, initial_time=m.t.at(4), representative_time=m.t.at(5)
     )
     t4 = m.t.at(4)
     for j in range(1, 6):
@@ -1456,8 +1518,7 @@ def test_get_initial_condition_problem_non_time_indexed_constraint_not_t0():
     m.eq_ydot3[t4].deactivate()
     m.eq_ydot4[t4].deactivate()
     m.eq_ydot5[t4].deactivate()
-    init_subsystem = petsc.get_initial_condition_problem(model=m, time=m.t, t_initial=m.t.at(4), representative_time=m.t.at(5))
+    init_subsystem = petsc.get_initial_condition_problem(model=m, time=m.t, initial_time=m.t.at(4), representative_time=m.t.at(5))
     from idaes.core.util.model_diagnostics import DiagnosticsToolbox
     dt=DiagnosticsToolbox(init_subsystem)
     dt.report_structural_issues()
-    print("ok, boomer")

From d54284afca71418c1a7004462fc2975e3e6d6327 Mon Sep 17 00:00:00 2001
From: Doug A <douglas.allan@keylogic.com>
Date: Wed, 26 Jun 2024 15:29:04 -0400
Subject: [PATCH 4/4] Run Black

---
 idaes/core/solvers/petsc.py            |  65 +++----
 idaes/core/solvers/tests/test_petsc.py | 230 ++++++++++++++++++++-----
 2 files changed, 222 insertions(+), 73 deletions(-)

diff --git a/idaes/core/solvers/petsc.py b/idaes/core/solvers/petsc.py
index 7129ccbee5..8cff031efb 100644
--- a/idaes/core/solvers/petsc.py
+++ b/idaes/core/solvers/petsc.py
@@ -42,7 +42,10 @@
 import idaes
 import idaes.logger as idaeslog
 import idaes.config as icfg
-from idaes.core.util.model_statistics import degrees_of_freedom, number_activated_constraints
+from idaes.core.util.model_statistics import (
+    degrees_of_freedom,
+    number_activated_constraints,
+)
 
 
 # Importing a few things here so that they are cached
@@ -252,7 +255,7 @@ def _copy_time(time_vars, t_from, t_to):
 def find_discretization_equations(m, time):
     """
     This returns a list of continuity equations (if Lagrange-Legendre collocation is used)
-    and time discretization equations. Since we aren't solving the whole time period 
+    and time discretization equations. Since we aren't solving the whole time period
     simultaneously, we'll want to deactivate these constraints.
 
     Args:
@@ -290,6 +293,7 @@ def find_discretization_equations(m, time):
 
     return disc_eqns
 
+
 def get_discretization_constraint_data(m, time):
     """
     This returns a list of continuity equations (if Lagrange-Legendre collocation is used)
@@ -567,10 +571,8 @@ def petsc_dae_by_time_element(
     elif representative_time not in between:
         raise RuntimeError("representative_time is not element of between.")
 
-    flattened_problem  = _get_flattened_problem(
-        model=m,
-        time=time,
-        representative_time=representative_time
+    flattened_problem = _get_flattened_problem(
+        model=m, time=time, representative_time=representative_time
     )
 
     solver_dae = pyo.SolverFactory("petsc_ts", options=ts_options)
@@ -604,10 +606,13 @@ def petsc_dae_by_time_element(
             )
         except RuntimeError as err:
             if "Zero constraints" in err.message:
-                init_subsystem=None
+                init_subsystem = None
             else:
                 raise err
-        if init_subsystem is not None and number_activated_constraints(init_subsystem) > 0:
+        if (
+            init_subsystem is not None
+            and number_activated_constraints(init_subsystem) > 0
+        ):
             dof = degrees_of_freedom(init_subsystem)
             if dof > 0:
                 raise RuntimeError(
@@ -639,7 +644,9 @@ def petsc_dae_by_time_element(
             if t == between.first():
                 # t == between.first() was handled above
                 continue
-            constraints = [con[t] for con in flattened_problem["time_constraints"] if t in con]
+            constraints = [
+                con[t] for con in flattened_problem["time_constraints"] if t in con
+            ]
             variables = [var[t] for var in flattened_problem["time_variables"]]
             # Create a temporary block with references to original constraints
             # and variables so we can integrate this "subsystem" without
@@ -1057,6 +1064,7 @@ def from_json(self, pth):
                 self.vecs = json.load(fp)
         self.time = self.vecs["_time"]
 
+
 def _get_flattened_problem(model, time, representative_time):
     """
     Helper function for petsc_dae_by_time_element and get_initial_condition_problem.
@@ -1067,11 +1075,11 @@ def _get_flattened_problem(model, time, representative_time):
         time (ContinuousSet): Time set
         representative_time (Element of time): A timepoint at which the DAE system is at its "normal" state
             after the constraints and variables associated with the initial time point
-            have passed. 
+            have passed.
 
     Returns (dictionary):
         Dictionary containing lists of variables and constraints indexed by times and lists
-        of those unindexed by time. Also a list of discretization equations so they can be 
+        of those unindexed by time. Also a list of discretization equations so they can be
         deactivated in the problem passed to PETSc.
     """
     regular_vars, time_vars = flatten_dae_components(
@@ -1092,15 +1100,15 @@ def _get_flattened_problem(model, time, representative_time):
 
 
 def get_initial_condition_problem(
-        model,
-        time,
-        initial_time,
-        representative_time=None,
-        initial_constraints=None,
-        initial_variables=None,
-        detect_initial=True,
-        flattened_problem=None,
-    ):
+    model,
+    time,
+    initial_time,
+    representative_time=None,
+    initial_constraints=None,
+    initial_variables=None,
+    detect_initial=True,
+    flattened_problem=None,
+):
     """
     Solve a DAE problem step by step using the PETSc DAE solver.  This
     integrates from one time point to the next.
@@ -1139,11 +1147,11 @@ def get_initial_condition_problem(
 
     if flattened_problem is None:
         if representative_time is None:
-            raise RuntimeError("The user must supply either the flattened problem or a representative time.")
+            raise RuntimeError(
+                "The user must supply either the flattened problem or a representative time."
+            )
         flattened_problem = _get_flattened_problem(
-            model=model,
-            time=time,
-            representative_time=representative_time
+            model=model, time=time, representative_time=representative_time
         )
 
     if detect_initial:
@@ -1156,17 +1164,14 @@ def get_initial_condition_problem(
         const_init_set = ComponentSet(initial_constraints)
         initial_constraints = list(const_no_t_set | const_init_set)
 
-
-
     constraints = [
-        con[initial_time] 
+        con[initial_time]
         for con in flattened_problem["time_constraints"]
         if initial_time in con
         and con[initial_time] not in flattened_problem["discretization_equations"]
     ] + initial_constraints
     variables = [
-        var[initial_time] 
-        for var in flattened_problem["time_variables"]
+        var[initial_time] for var in flattened_problem["time_variables"]
     ] + initial_variables
 
     if len(constraints) <= 0:
@@ -1181,4 +1186,4 @@ def get_initial_condition_problem(
         variables,
     )
     _sub_problem_scaling_suffix(model, subsystem)
-    return subsystem
\ No newline at end of file
+    return subsystem
diff --git a/idaes/core/solvers/tests/test_petsc.py b/idaes/core/solvers/tests/test_petsc.py
index 846ef067c4..13ac59531e 100644
--- a/idaes/core/solvers/tests/test_petsc.py
+++ b/idaes/core/solvers/tests/test_petsc.py
@@ -812,6 +812,7 @@ def test_petsc_skip_initial_solve():
     assert pytest.approx(y5, rel=1e-3) == pyo.value(m.y[m.t.last(), 5])
     assert pytest.approx(y6, rel=1e-3) == pyo.value(m.y[m.t.last(), 6])
 
+
 @pytest.mark.unit
 @pytest.mark.skipif(not petsc.petsc_available(), reason="PETSc solver not available")
 def test_petsc_collocation_lagrange_radau():
@@ -819,9 +820,7 @@ def test_petsc_collocation_lagrange_radau():
     Ensure that PETSc works with Lagrange-Radau collocation
     """
     m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(
-        transformation_method="dae.collocation",
-        scheme="LAGRANGE-RADAU",
-        ncp=3
+        transformation_method="dae.collocation", scheme="LAGRANGE-RADAU", ncp=3
     )
 
     res = petsc.petsc_dae_by_time_element(
@@ -843,6 +842,7 @@ def test_petsc_collocation_lagrange_radau():
     assert pytest.approx(y5, rel=1e-3) == pyo.value(m.y[m.t.last(), 5])
     assert pytest.approx(y6, rel=1e-3) == pyo.value(m.y[m.t.last(), 6])
 
+
 @pytest.mark.unit
 @pytest.mark.skipif(not petsc.petsc_available(), reason="PETSc solver not available")
 def test_petsc_collocation_lagrange_legendre():
@@ -850,9 +850,7 @@ def test_petsc_collocation_lagrange_legendre():
     Ensure that PETSc works with Lagrange-Legendre collocation
     """
     m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(
-        transformation_method="dae.collocation",
-        scheme="LAGRANGE-LEGENDRE",
-        ncp=3
+        transformation_method="dae.collocation", scheme="LAGRANGE-LEGENDRE", ncp=3
     )
 
     res = petsc.petsc_dae_by_time_element(
@@ -932,11 +930,11 @@ def test_petsc_traj_previous():
         initial_solver="ipopt",
         initial_solver_options={
             # 'bound_push' : 1e-22,
-            'linear_solver': 'ma57',
-            'OF_ma57_automatic_scaling': 'yes',
-            'max_iter': 300,
-            'tol': 1e-6,
-            'halt_on_ampl_error': 'yes',
+            "linear_solver": "ma57",
+            "OF_ma57_automatic_scaling": "yes",
+            "max_iter": 300,
+            "tol": 1e-6,
+            "halt_on_ampl_error": "yes",
             # 'mu_strategy': 'monotone',
         },
         # skip_initial=True,
@@ -1353,6 +1351,7 @@ def diff_eq_rule(m, t):
             calculate_derivatives=True,
         )
 
+
 @pytest.mark.unit
 def test_get_initial_condition_problem_no_active_constraints():
     m = pyo.ConcreteModel()
@@ -1378,16 +1377,21 @@ def diff_eq_rule(m, t):
         model=m,
         time=m.time,
         initial_time=m.time.at(1),
-        representative_time=m.time.at(2)
+        representative_time=m.time.at(2),
     )
     assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
-    local_var_name_list = [var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)]
-    assert set(local_var_name_list) == {'x[0.0]', 'u[0.0]', 'dxdt[0.0]'}
+    local_var_name_list = [
+        var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)
+    ]
+    assert set(local_var_name_list) == {"x[0.0]", "u[0.0]", "dxdt[0.0]"}
 
     assert mstat.number_activated_constraints(init_subsystem) == 0
-    local_con_name_list = [con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)]
+    local_con_name_list = [
+        con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)
+    ]
     assert local_con_name_list[0] == "diff_eq[0.0]"
 
+
 @pytest.mark.unit
 def test_get_initial_condition_problem_no_active_constraints_not_t0():
     m = pyo.ConcreteModel()
@@ -1413,16 +1417,21 @@ def diff_eq_rule(m, t):
         model=m,
         time=m.time,
         initial_time=m.time.at(2),
-        representative_time=m.time.at(3)
+        representative_time=m.time.at(3),
     )
     assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
-    local_var_name_list = [var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)]
-    assert set(local_var_name_list) == {'x[1.0]', 'u[1.0]', 'dxdt[1.0]'}
+    local_var_name_list = [
+        var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)
+    ]
+    assert set(local_var_name_list) == {"x[1.0]", "u[1.0]", "dxdt[1.0]"}
 
     assert mstat.number_activated_constraints(init_subsystem) == 0
-    local_con_name_list = [con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)]
+    local_con_name_list = [
+        con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)
+    ]
     assert local_con_name_list[0] == "diff_eq[1.0]"
 
+
 @pytest.mark.unit
 def test_get_initial_condition_problem_at_time_no_active_constraints():
     m = pyo.ConcreteModel()
@@ -1445,34 +1454,103 @@ def diff_eq_rule(m, t):
         model=m,
         time=m.time,
         initial_time=m.time.at(1),
-        representative_time=m.time.at(2)
+        representative_time=m.time.at(2),
     )
     assert mstat.number_variables_in_activated_constraints(init_subsystem) == 0
-    local_var_name_list = [var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)]
-    assert set(local_var_name_list) == {'x[0.0]', 'u[0.0]', 'dxdt[0.0]'}
+    local_var_name_list = [
+        var.name for var in init_subsystem.component_data_objects(ctype=pyo.Var)
+    ]
+    assert set(local_var_name_list) == {"x[0.0]", "u[0.0]", "dxdt[0.0]"}
 
     assert mstat.number_activated_constraints(init_subsystem) == 0
-    local_con_name_list = [con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)]
+    local_con_name_list = [
+        con.name for con in init_subsystem.component_data_objects(ctype=pyo.Constraint)
+    ]
     assert local_con_name_list[0] == "diff_eq[0.0]"
 
+
 @pytest.mark.unit
 def test_get_initial_condition_problem_non_time_indexed_constraint():
     m, y1, y2, y3, y4, y5, y6 = dae_with_non_time_indexed_constraint(nfe=10)
     init_subsystem = petsc.get_initial_condition_problem(
         model=m, time=m.t, initial_time=m.t.at(1), representative_time=m.t.at(2)
     )
-    
-    unfixed_var_name_set = set(var.name for var in mstat.unfixed_variables_set(init_subsystem))
-    assert unfixed_var_name_set == {'y[0,6]', 'ydot[0,1]', 'ydot[0,2]', 'ydot[0,3]', 'ydot[0,4]', 'ydot[0,5]', 'r[0,1]', 'r[0,2]', 'r[0,3]', 'r[0,4]', 'r[0,5]', 'Fin[0]', 'H'}
 
-    total_var_name_set = set(var.name for var in mstat.variables_set(init_subsystem))
-    assert total_var_name_set == {'r[0,1]', 'r[0,2]', 'ydot[0,1]', 'ydot[0,4]', 'H', 'y[0,2]', 'ydot[0,3]', 'y[0,3]', 'y[0,5]', 'r[0,3]', 'r[0,5]', 'Fin[0]', 'y[0,1]', 'ydot[0,2]', 'ydot[0,6]', 'r[0,4]', 'y[0,4]', 'y[0,6]', 'ydot[0,5]'}
+    unfixed_var_name_set = set(
+        var.name for var in mstat.unfixed_variables_set(init_subsystem)
+    )
+    assert unfixed_var_name_set == {
+        "y[0,6]",
+        "ydot[0,1]",
+        "ydot[0,2]",
+        "ydot[0,3]",
+        "ydot[0,4]",
+        "ydot[0,5]",
+        "r[0,1]",
+        "r[0,2]",
+        "r[0,3]",
+        "r[0,4]",
+        "r[0,5]",
+        "Fin[0]",
+        "H",
+    }
 
-    active_con_name_set = set(con.name for con in mstat.activated_constraints_set(init_subsystem))
-    assert active_con_name_set == {'eq_y6[0]', 'eq_r1[0]', 'eq_r2[0]', 'eq_r3[0]', 'eq_r4[0]', 'eq_r5[0]', 'eq_Fin[0]', 'H_eqn'}
+    total_var_name_set = set(var.name for var in mstat.variables_set(init_subsystem))
+    assert total_var_name_set == {
+        "r[0,1]",
+        "r[0,2]",
+        "ydot[0,1]",
+        "ydot[0,4]",
+        "H",
+        "y[0,2]",
+        "ydot[0,3]",
+        "y[0,3]",
+        "y[0,5]",
+        "r[0,3]",
+        "r[0,5]",
+        "Fin[0]",
+        "y[0,1]",
+        "ydot[0,2]",
+        "ydot[0,6]",
+        "r[0,4]",
+        "y[0,4]",
+        "y[0,6]",
+        "ydot[0,5]",
+    }
+
+    active_con_name_set = set(
+        con.name for con in mstat.activated_constraints_set(init_subsystem)
+    )
+    assert active_con_name_set == {
+        "eq_y6[0]",
+        "eq_r1[0]",
+        "eq_r2[0]",
+        "eq_r3[0]",
+        "eq_r4[0]",
+        "eq_r5[0]",
+        "eq_Fin[0]",
+        "H_eqn",
+    }
+
+    total_con_name_set = set(
+        con.name for con in mstat.total_constraints_set(init_subsystem)
+    )
+    assert total_con_name_set == {
+        "eq_r5[0]",
+        "eq_y6[0]",
+        "eq_r4[0]",
+        "eq_Fin[0]",
+        "H_eqn",
+        "eq_ydot5[0]",
+        "eq_ydot1[0]",
+        "eq_r1[0]",
+        "eq_ydot4[0]",
+        "eq_r3[0]",
+        "eq_r2[0]",
+        "eq_ydot2[0]",
+        "eq_ydot3[0]",
+    }
 
-    total_con_name_set = set(con.name for con in mstat.total_constraints_set(init_subsystem))
-    assert total_con_name_set == {'eq_r5[0]', 'eq_y6[0]', 'eq_r4[0]', 'eq_Fin[0]', 'H_eqn', 'eq_ydot5[0]', 'eq_ydot1[0]', 'eq_r1[0]', 'eq_ydot4[0]', 'eq_r3[0]', 'eq_r2[0]', 'eq_ydot2[0]', 'eq_ydot3[0]'}
 
 @pytest.mark.unit
 def test_get_initial_condition_problem_non_time_indexed_constraint_not_t0():
@@ -1489,17 +1567,80 @@ def test_get_initial_condition_problem_non_time_indexed_constraint_not_t0():
     m.eq_ydot4[t4].deactivate()
     m.eq_ydot5[t4].deactivate()
 
-    unfixed_var_name_set = set(var.name for var in mstat.unfixed_variables_set(init_subsystem))
-    assert unfixed_var_name_set == {'y[54.0,6]', 'ydot[54.0,1]', 'ydot[54.0,2]', 'ydot[54.0,3]', 'ydot[54.0,4]', 'ydot[54.0,5]', 'r[54.0,1]', 'r[54.0,2]', 'r[54.0,3]', 'r[54.0,4]', 'r[54.0,5]', 'Fin[54.0]', 'H'}
+    unfixed_var_name_set = set(
+        var.name for var in mstat.unfixed_variables_set(init_subsystem)
+    )
+    assert unfixed_var_name_set == {
+        "y[54.0,6]",
+        "ydot[54.0,1]",
+        "ydot[54.0,2]",
+        "ydot[54.0,3]",
+        "ydot[54.0,4]",
+        "ydot[54.0,5]",
+        "r[54.0,1]",
+        "r[54.0,2]",
+        "r[54.0,3]",
+        "r[54.0,4]",
+        "r[54.0,5]",
+        "Fin[54.0]",
+        "H",
+    }
 
     total_var_name_set = set(var.name for var in mstat.variables_set(init_subsystem))
-    assert total_var_name_set == {'r[54.0,1]', 'r[54.0,2]', 'ydot[54.0,1]', 'ydot[54.0,4]', 'H', 'y[54.0,2]', 'ydot[54.0,3]', 'y[54.0,3]', 'y[54.0,5]', 'r[54.0,3]', 'r[54.0,5]', 'Fin[54.0]', 'y[54.0,1]', 'ydot[54.0,2]', 'ydot[54.0,6]', 'r[54.0,4]', 'y[54.0,4]', 'y[54.0,6]', 'ydot[54.0,5]'}
-
-    active_con_name_set = set(con.name for con in mstat.activated_constraints_set(init_subsystem))
-    assert active_con_name_set == {'eq_y6[54.0]', 'eq_r1[54.0]', 'eq_r2[54.0]', 'eq_r3[54.0]', 'eq_r4[54.0]', 'eq_r5[54.0]', 'eq_Fin[54.0]', 'H_eqn'}
-
-    total_con_name_set = set(con.name for con in mstat.total_constraints_set(init_subsystem))
-    assert total_con_name_set == {'eq_r5[54.0]', 'eq_y6[54.0]', 'eq_r4[54.0]', 'eq_Fin[54.0]', 'H_eqn', 'eq_ydot5[54.0]', 'eq_ydot1[54.0]', 'eq_r1[54.0]', 'eq_ydot4[54.0]', 'eq_r3[54.0]', 'eq_r2[54.0]', 'eq_ydot2[54.0]', 'eq_ydot3[54.0]'}
+    assert total_var_name_set == {
+        "r[54.0,1]",
+        "r[54.0,2]",
+        "ydot[54.0,1]",
+        "ydot[54.0,4]",
+        "H",
+        "y[54.0,2]",
+        "ydot[54.0,3]",
+        "y[54.0,3]",
+        "y[54.0,5]",
+        "r[54.0,3]",
+        "r[54.0,5]",
+        "Fin[54.0]",
+        "y[54.0,1]",
+        "ydot[54.0,2]",
+        "ydot[54.0,6]",
+        "r[54.0,4]",
+        "y[54.0,4]",
+        "y[54.0,6]",
+        "ydot[54.0,5]",
+    }
+
+    active_con_name_set = set(
+        con.name for con in mstat.activated_constraints_set(init_subsystem)
+    )
+    assert active_con_name_set == {
+        "eq_y6[54.0]",
+        "eq_r1[54.0]",
+        "eq_r2[54.0]",
+        "eq_r3[54.0]",
+        "eq_r4[54.0]",
+        "eq_r5[54.0]",
+        "eq_Fin[54.0]",
+        "H_eqn",
+    }
+
+    total_con_name_set = set(
+        con.name for con in mstat.total_constraints_set(init_subsystem)
+    )
+    assert total_con_name_set == {
+        "eq_r5[54.0]",
+        "eq_y6[54.0]",
+        "eq_r4[54.0]",
+        "eq_Fin[54.0]",
+        "H_eqn",
+        "eq_ydot5[54.0]",
+        "eq_ydot1[54.0]",
+        "eq_r1[54.0]",
+        "eq_ydot4[54.0]",
+        "eq_r3[54.0]",
+        "eq_r2[54.0]",
+        "eq_ydot2[54.0]",
+        "eq_ydot3[54.0]",
+    }
 
 
 if __name__ == "__main__":
@@ -1518,7 +1659,10 @@ def test_get_initial_condition_problem_non_time_indexed_constraint_not_t0():
     m.eq_ydot3[t4].deactivate()
     m.eq_ydot4[t4].deactivate()
     m.eq_ydot5[t4].deactivate()
-    init_subsystem = petsc.get_initial_condition_problem(model=m, time=m.t, initial_time=m.t.at(4), representative_time=m.t.at(5))
+    init_subsystem = petsc.get_initial_condition_problem(
+        model=m, time=m.t, initial_time=m.t.at(4), representative_time=m.t.at(5)
+    )
     from idaes.core.util.model_diagnostics import DiagnosticsToolbox
-    dt=DiagnosticsToolbox(init_subsystem)
+
+    dt = DiagnosticsToolbox(init_subsystem)
     dt.report_structural_issues()