Skip to content

Commit

Permalink
control_comp now handles constraints on control continuity and rate c…
Browse files Browse the repository at this point in the history 
…ontinuity
  • Loading branch information
@robfalck
robfalck committed Nov 18, 2024
1 parent 06b5ea1 commit ec57b5d
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 118 deletions.
86 changes: 2 additions & 84 deletions dymos/transcriptions/common/control_group.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,9 +68,6 @@ def initialize(self):
self._output_val_names = {}
self._output_rate_names = {}
self._output_rate2_names = {}
self._output_cnty_defect_names = {}
self._output_rate_cnty_defect_names = {}
self._output_rate2_cnty_defect_names = {}
self._matrices = {}

def setup(self):
Expand All @@ -88,8 +85,6 @@ def setup(self):

def _configure_controls(self):
gd = self.options['grid_data']
num_seg = gd.num_segments
segment_end_idxs = gd.subset_node_indices['segment_ends']
ogd = self.options['output_grid_data'] or gd
eval_nodes = ogd.node_ptau
control_options = self.options['control_options']
Expand Down Expand Up @@ -181,16 +176,10 @@ def _configure_controls(self):
self._output_val_names[name] = f'control_values:{name}'
self._output_rate_names[name] = f'control_rates:{name}_rate'
self._output_rate2_names[name] = f'control_rates:{name}_rate2'
self._output_cnty_defect_names[name] = f'cnty_defect_controls:{name}'
self._output_rate_cnty_defect_names[name] = f'cnty_defect_control_rates:{name}_rate'
self._output_rate2_cnty_defect_names[name] = f'cnty_defect_control_rates:{name}_rate2'

shape = options['shape']
num_control_input_nodes = gd.subset_num_nodes['control_input']
ncinps = gd.subset_num_nodes_per_segment['control_input']
input_shape = (num_control_input_nodes,) + shape
output_shape = (num_output_nodes,) + shape
cnty_output_shape = (num_seg - 1,) + shape

units = options['units']
rate_units = get_rate_units(units, time_units)
Expand All @@ -205,13 +194,6 @@ def _configure_controls(self):
self.add_output(self._output_rate2_names[name], shape=output_shape,
units=rate2_units)

self.add_output(self._output_cnty_defect_names[name], shape=cnty_output_shape, units=units)

self.add_output(self._output_rate_cnty_defect_names[name], shape=cnty_output_shape, units=rate_units)

self.add_output(self._output_rate2_cnty_defect_names[name], shape=cnty_output_shape,
units=rate2_units)

size = np.prod(shape)
self.sizes[name] = size
sp_eye = sp.eye(size, format='csr')
Expand All @@ -225,17 +207,6 @@ def _configure_controls(self):
wrt=self._input_names[name],
rows=rs, cols=cs, val=data)

# The jacobian for continuity defects is similar to the jacobian of the inteprpolated values
J_cnty_def = -J_val[segment_end_idxs, ...][1:-1:2, ...].copy()
one_col_idxs = np.cumsum(np.asarray(ncinps[:-1]))
J_cnty_def[np.arange(num_seg -1, dtype=int), one_col_idxs] = 1.0

rs, cs, data = sp.find(J_cnty_def)

self.declare_partials(of=self._output_cnty_defect_names[name],
wrt=self._input_names[name],
rows=rs, cols=cs, val=data)

# The partials of the output rate and second derivative wrt dt_dstau
rs = np.arange(num_output_nodes * size, dtype=int)
cs = np.repeat(np.arange(num_output_nodes, dtype=int), size)
Expand All @@ -258,31 +229,7 @@ def _configure_controls(self):
wrt=self._input_names[name],
rows=rs, cols=cs)

# # Similar with the continuity rate defects.
J_rate_cnty_def = sp.kron(self.D, sp_eye, format='csr')[segment_end_idxs, ...].copy()
J_rate_cnty_def = J_rate_cnty_def[1:-1]
J_rate_cnty_def[::2] *= -1
J_rate_cnty_def /= gd.node_dptau_dstau[segment_end_idxs, np.newaxis][1:-1]

# with np.printoptions(linewidth=10000):
# print(J_rate_cnty_def.todense())
# rs, cs, data = sp.find(J_rate_cnty_def)
# print(rs)
# print(gd.subset_num_nodes_per_segment['control_input'])
# print(rs - [0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 4, 4, 4])
# # compute the row repeats
# for i in range(num_seg):
# print(np.arange(num_seg, dtype=int))
# print(np.asarray(gd.subset_num_nodes_per_segment['control_input'], dtype=int) * 2)
# # print(cs)
# exit(0)

# self.declare_partials(of=self._output_rate_cnty_defect_names[name],
# wrt=self._input_names[name],
# rows=rs, cols=cs, val=data)

self.rate2_jacs[name] = sp.kron(self.D2, sp_eye, format='csr')

rs, cs = self.rate2_jacs[name].nonzero()

self.declare_partials(of=self._output_rate2_names[name],
Expand Down Expand Up @@ -380,9 +327,7 @@ def compute(self, inputs, outputs):
control_options = self.options['control_options']
num_control_input_nodes = self.options['grid_data'].subset_num_nodes['control_input']
num_output_nodes = ogd.num_nodes
seg_end_idxs = gd.subset_node_indices['segment_ends']
dt_dptau = 0.5 * inputs['t_duration']
dptau_dstau = gd.node_dptau_dstau

for name, options in control_options.items():
if control_options[name]['control_type'] == 'polynomial':
Expand All @@ -403,8 +348,8 @@ def compute(self, inputs, outputs):
u_flat = np.reshape(inputs[self._input_names[name]],
(num_control_input_nodes, size))

a = self.D.dot(u_flat) # du/dstau
b = self.D2.dot(u_flat) # d2u/dstau2
a = self.D.dot(u_flat)
b = self.D2.dot(u_flat)

val = np.reshape(self.L.dot(u_flat), (num_output_nodes,) + options['shape'])

Expand All @@ -418,33 +363,6 @@ def compute(self, inputs, outputs):
outputs[self._output_rate_names[name]] = rate
outputs[self._output_rate2_names[name]] = rate2

if True:#options['continuity']:
# output_name = self._output_val_names[name]
cnty_output_name = self._output_cnty_defect_names[name]
segment_end_values = val[seg_end_idxs, ...]
end_vals = segment_end_values[1:-1:2, ...]
start_vals = segment_end_values[2:-1:2, ...]
outputs[cnty_output_name] = start_vals - end_vals

if True:#options['rate_continuity']:
# output_name = self._output_rate_names[name]
rate_in_ptau = a / dptau_dstau[:, np.newaxis]
rate_in_ptau = np.reshape(rate_in_ptau, (num_output_nodes,) + options['shape'])

cnty_output_name = self._output_rate_cnty_defect_names[name]
segment_end_values = rate_in_ptau[seg_end_idxs, ...]
end_vals = segment_end_values[1:-1:2, ...]
start_vals = segment_end_values[2:-1:2, ...]
outputs[cnty_output_name] = (start_vals - end_vals) #* dt_dptau

if True:#options['rate2_continuity']:
# output_name = self._output_rate2_names[name]
cnty_output_name = self._output_rate2_cnty_defect_names[name]
segment_end_values = rate2[seg_end_idxs, ...]
end_vals = segment_end_values[1:-1:2, ...]
start_vals = segment_end_values[2:-1:2, ...]
outputs[cnty_output_name] = (start_vals - end_vals) #* dt_dptau ** 2

def compute_partials(self, inputs, partials):
"""
Compute sub-jacobian parts. The model is assumed to be in an unscaled state.
Expand Down
70 changes: 36 additions & 34 deletions dymos/transcriptions/pseudospectral/radau_new.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import openmdao.api as om

from ..transcription_base import TranscriptionBase
from ..common import TimeComp, TimeseriesOutputGroup, TimeseriesOutputComp, RadauPSContinuityComp, ControlComp
from ..common import TimeComp, TimeseriesOutputGroup, TimeseriesOutputComp, ControlComp
from .components import RadauIterGroup, RadauBoundaryGroup

from ..grid_data import RadauGrid
Expand Down Expand Up @@ -269,13 +269,14 @@ def setup_defects(self, phase):
phase : dymos.Phase
The phase object to which this transcription instance applies.
"""
if any(self._requires_continuity_constraints(phase)):
phase.add_subsystem('continuity_comp',
RadauPSContinuityComp(grid_data=self.grid_data,
# State continuity defects are in the RadauDefectComp.
state_options={},
control_options=phase.control_options,
time_units=phase.time_options['units']))
pass
# if any(self._requires_continuity_constraints(phase)):
# phase.add_subsystem('continuity_comp',
# RadauPSContinuityComp(grid_data=self.grid_data,
# # State continuity defects are in the RadauDefectComp.
# state_options={},
# control_options=phase.control_options,
# time_units=phase.time_options['units']))

def configure_defects(self, phase):
"""
Expand All @@ -295,41 +296,41 @@ def configure_defects(self, phase):
phase.connect(rate_src_path, f'f_ode:{name}',
src_indices=grid_data.subset_node_indices['col'])

any_state_cnty, any_control_cnty, any_control_rate_cnty = self._requires_continuity_constraints(phase)
# any_state_cnty, any_control_cnty, any_control_rate_cnty = self._requires_continuity_constraints(phase)

if not any((any_state_cnty, any_control_cnty, any_control_rate_cnty)):
return
# if not any((any_state_cnty, any_control_cnty, any_control_rate_cnty)):
# return

# Add the continuity constraint component if necessary
segment_end_idxs = grid_data.subset_node_indices['segment_ends']
# # Add the continuity constraint component if necessary
# segment_end_idxs = grid_data.subset_node_indices['segment_ends']

if any_control_rate_cnty:
phase.connect('t_duration_val', 'continuity_comp.t_duration')
# if any_control_rate_cnty:
# phase.connect('t_duration_val', 'continuity_comp.t_duration')

phase.continuity_comp.configure_io()
# phase.continuity_comp.configure_io()

for name, options in phase.control_options.items():
# The sub-indices of control_disc indices that are segment ends
segment_end_idxs = grid_data.subset_node_indices['segment_ends']
src_idxs = get_src_indices_by_row(segment_end_idxs, options['shape'], flat=True)
# for name, options in phase.control_options.items():
# # The sub-indices of control_disc indices that are segment ends
# segment_end_idxs = grid_data.subset_node_indices['segment_ends']
# src_idxs = get_src_indices_by_row(segment_end_idxs, options['shape'], flat=True)

# enclose indices in tuple to ensure shaping of indices works
src_idxs = (src_idxs,)
# # enclose indices in tuple to ensure shaping of indices works
# src_idxs = (src_idxs,)

if options['continuity']:
phase.connect(f'control_values:{name}',
f'continuity_comp.controls:{name}',
src_indices=src_idxs, flat_src_indices=True)
# if options['continuity']:
# phase.connect(f'control_values:{name}',
# f'continuity_comp.controls:{name}',
# src_indices=src_idxs, flat_src_indices=True)

if options['rate_continuity']:
phase.connect(f'control_rates:{name}_rate',
f'continuity_comp.control_rates:{name}_rate',
src_indices=src_idxs, flat_src_indices=True)
# if options['rate_continuity']:
# phase.connect(f'control_rates:{name}_rate',
# f'continuity_comp.control_rates:{name}_rate',
# src_indices=src_idxs, flat_src_indices=True)

if options['rate2_continuity']:
phase.connect(f'control_rates:{name}_rate2',
f'continuity_comp.control_rates:{name}_rate2',
src_indices=src_idxs, flat_src_indices=True)
# if options['rate2_continuity']:
# phase.connect(f'control_rates:{name}_rate2',
# f'continuity_comp.control_rates:{name}_rate2',
# src_indices=src_idxs, flat_src_indices=True)

def setup_duration_balance(self, phase):
"""
Expand Down Expand Up @@ -845,6 +846,7 @@ def get_parameter_connections(self, name, phase):
for tgt in options['targets']:
if tgt in options['static_targets']:
src_idxs = np.squeeze(get_src_indices_by_row([0], options['shape']), axis=0)
endpoint_src_idxs = om.slicer[:, ...]
else:
src_idxs_raw = np.zeros(self.grid_data.subset_num_nodes['all'], dtype=int)
src_idxs = get_src_indices_by_row(src_idxs_raw, options['shape'])
Expand Down

0 comments on commit ec57b5d

Please sign in to comment.