OpenMDAO / dymos

Open Source Optimization of Dynamic Multidisciplinary Systems
Apache License 2.0
202 stars 65 forks source link

Solver-based shooting with a Krylov solver fails due to the DirectSolver attached to `indep_states` #1009

Closed kanekosh closed 9 months ago

kanekosh commented 10 months ago

Description

When using Newton solver-based shooting with a Krylov linear solver, the LinearRunOnce preconditioner fails with the following error. It seems the error comes from the DirectSolver attached to indep_states component.

Dymos 1.9.0 didn't cause this error because indep_states didn't have a DirectSolver for the shooting settings. I believe the behavior changed since PR #991 .

Error message

Traceback (most recent call last):
  File "/Users/shugo/dymos_bug/brachistochrone_shooting.py", line 77, in <module>
    dm.run_problem(p, run_driver=False)
  File "/Users/shugo/packages/dymos/dymos/run_problem.py", line 100, in run_problem
    failed = problem.run_model()
             ^^^^^^^^^^^^^^^^^^^
  File "/Users/shugo/packages/openmdao-dev/openmdao/core/problem.py", line 654, in run_model
    self.model.run_solve_nonlinear()
  File "/Users/shugo/packages/openmdao-dev/openmdao/core/system.py", line 4552, in run_solve_nonlinear
    self._solve_nonlinear()
  File "/Users/shugo/packages/openmdao-dev/openmdao/core/group.py", line 3593, in _solve_nonlinear
    self._nonlinear_solver._solve_with_cache_check()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/nonlinear/nonlinear_runonce.py", line 26, in _solve_with_cache_check
    self.solve()  # don't use caching
    ^^^^^^^^^^^^
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/nonlinear/nonlinear_runonce.py", line 45, in solve
    self._gs_iter()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/solver.py", line 803, in _gs_iter
    subsys._solve_nonlinear()
  File "/Users/shugo/packages/openmdao-dev/openmdao/core/group.py", line 3593, in _solve_nonlinear
    self._nonlinear_solver._solve_with_cache_check()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/nonlinear/nonlinear_runonce.py", line 26, in _solve_with_cache_check
    self.solve()  # don't use caching
    ^^^^^^^^^^^^
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/nonlinear/nonlinear_runonce.py", line 45, in solve
    self._gs_iter()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/solver.py", line 803, in _gs_iter
    subsys._solve_nonlinear()
  File "/Users/shugo/packages/openmdao-dev/openmdao/core/group.py", line 3593, in _solve_nonlinear
    self._nonlinear_solver._solve_with_cache_check()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/nonlinear/nonlinear_runonce.py", line 26, in _solve_with_cache_check
    self.solve()  # don't use caching
    ^^^^^^^^^^^^
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/nonlinear/nonlinear_runonce.py", line 45, in solve
    self._gs_iter()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/solver.py", line 803, in _gs_iter
    subsys._solve_nonlinear()
  File "/Users/shugo/packages/openmdao-dev/openmdao/core/group.py", line 3593, in _solve_nonlinear
    self._nonlinear_solver._solve_with_cache_check()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/solver.py", line 845, in _solve_with_cache_check
    self.solve()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/solver.py", line 592, in solve
    raise err
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/solver.py", line 588, in solve
    self._solve()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/solver.py", line 674, in _solve
    self._single_iteration()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/nonlinear/newton.py", line 227, in _single_iteration
    system._linearize(my_asm_jac, sub_do_ln=do_sub_ln)
  File "/Users/shugo/packages/openmdao-dev/openmdao/core/group.py", line 3822, in _linearize
    subsys._linear_solver._linearize()
  File "/Users/shugo/packages/openmdao-dev/openmdao/solvers/linear/direct.py", line 277, in _linearize
    matrix = self._assembled_jac._int_mtx._matrix
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute '_matrix'

Workarounds

Potential fix

I believe attaching aDirectSolver for indep_states is not necessary for shooting cases, because Dymos has a DirectSolver at the phase level anyway.

Example

Here is a Brachistochrone example. I made a few changes to the collocation example in the Dymos docs: 1. Use solver-based shooting

phase0 = dm.Phase(ode_class=BrachistochroneODE, transcription=dm.GaussLobatto(num_segments=10, compressed=True, solve_segments='forward'))
phase = traj.add_phase('phase0', phase0)

2. Change the phase-level linear solver to PETScKrylov

phase0.linear_solver = om.PETScKrylov(assemble_jac=True, iprint=2)
phase0.linear_solver.precon = om.LinearRunOnce(iprint=-1)

Complete script

import openmdao.api as om
import dymos as dm
from dymos.examples.plotting import plot_results
from dymos.examples.brachistochrone import BrachistochroneODE
import matplotlib.pyplot as plt

#
# Initialize the Problem and the optimization driver
#
p = om.Problem(model=om.Group())
# p.driver = om.ScipyOptimizeDriver()
# p.driver.declare_coloring()

#
# Create a trajectory and add a phase to it
#
traj = p.model.add_subsystem('traj', dm.Trajectory())

# -------------------------------
# forward solver-based shooting
# -------------------------------
phase0 = dm.Phase(ode_class=BrachistochroneODE, transcription=dm.GaussLobatto(num_segments=10, compressed=True, solve_segments='forward'))
phase = traj.add_phase('phase0', phase0)

#
# Set the variables
#
phase.set_time_options(fix_initial=True, duration_bounds=(.5, 10))

phase.add_state('x', fix_initial=True, fix_final=False)

phase.add_state('y', fix_initial=True, fix_final=False)

phase.add_state('v', fix_initial=True, fix_final=False)

phase.add_control('theta', continuity=True, rate_continuity=True,
                  units='deg', lower=0.01, upper=179.9)

phase.add_parameter('g', units='m/s**2', val=9.80665)

#
# Minimize time at the end of the phase
#
phase.add_objective('time', loc='final', scaler=10)

# ---------------------------------------------
# change linear solver for shooting
# ---------------------------------------------
# linear solver
phase0.linear_solver = om.PETScKrylov(assemble_jac=True, iprint=2)
phase0.linear_solver.precon = om.LinearRunOnce(iprint=-1)

#
# Setup the Problem
#
p.setup()

# ---------------------------------------------
# Remove DirectSolver that causes an error
# ---------------------------------------------
# traj.phases.phase0.indep_states.linear_solver = None

#
# Set the initial values
#
p['traj.phase0.t_initial'] = 0.0
p['traj.phase0.t_duration'] = 2.0

p.set_val('traj.phase0.states:x', phase.interp('x', ys=[0, 10]))
p.set_val('traj.phase0.states:y', phase.interp('y', ys=[10, 5]))
p.set_val('traj.phase0.states:v', phase.interp('v', ys=[0, 9.9]))
p.set_val('traj.phase0.controls:theta', phase.interp('theta', ys=[5, 100.5]))

#
# Solve for the optimal trajectory
#
dm.run_problem(p, run_driver=False)

# plot results
sol = om.CaseReader('dymos_solution.db').get_case('final')
plot_results([('traj.phase0.timeseries.x', 'traj.phase0.timeseries.y',
               'x', 'y'),
              ('traj.phase0.timeseries.x', 'traj.phase0.timeseries.v',
               'x', 'v')],
             title='plot',
             p_sol=sol, p_sim=None)
plt.show()

Dymos Version

1.9.1

Relevant environment information

OpenMDAO 3.28.0

robfalck commented 10 months ago

I'm thinking the correct way to handle this is to noisily add a Newton and Direct solver when appropriate by default, but to let the user know how to easily replace these with the PetscKrylov solver.