def compute(self, inputs, outputs):
num_times = self.options['num_times']
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
y0_name = get_name('y0', state_name)
F_name = get_name('F', state_name)
y_name = get_name('y', state_name)
mtx_lu = self.mtx_lu_dict[state_name]
mtx_y0 = self.mtx_y0_dict[state_name]
mtx_h = self.mtx_h_dict[state_name]
mtx_hf = self.mtx_hf_dict[state_name]
# ------------------------------------------------------------------------------
vec = mtx_h.dot(inputs['h_vec']) * inputs[F_name].flatten()
vec = mtx_hf.dot(vec)
vec = mtx_lu.solve(vec)
outputs[y_name] = vec.reshape((
num_times,
num_step_vars,
) + shape)
outputs[y_name][0, :, :] -= inputs[y0_name]
def compute(self, inputs, outputs):
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
glm_B = self.options['glm_B']
glm_V = self.options['glm_V']
i_step = self.options['i_step']
ii_step = 0
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
y_old_name = get_name('y_old', state_name, i_step=i_step)
y_new_name = get_name('y_new', state_name, i_step=i_step)
outputs[y_new_name] = np.einsum('ij,jk...->ik...', glm_V,
inputs[y_old_name])
for j_stage in range(num_stages):
F_name = get_name('F',
state_name,
i_step=i_step,
j_stage=j_stage)
#y_new:step x shape
# GLM step x stage
# F: 1 x shape
outputs[y_new_name] += inputs['h'] * np.einsum(
'i,...->i...', glm_B[:, j_stage], inputs[F_name][0, :])
def compute(self, inputs, outputs):
num_starting_times = self.options['num_starting_times']
num_my_times = self.options['num_my_times']
num_step_vars = self.options['num_step_vars']
starting_coeffs = self.options['starting_coeffs']
has_starting_method = num_starting_times > 1
is_starting_method = starting_coeffs is not None
for state_name, state in iteritems(self.options['states']):
starting_state_name = get_name('starting_state', state_name)
out_state_name = get_name('state', state_name)
starting_name = get_name('starting', state_name)
if has_starting_method:
outputs[out_state_name][:num_starting_times - 1] = \
inputs[starting_state_name][:-1, :]
if is_starting_method:
outputs[starting_name] = 0.
for i_step in range(num_my_times):
y_name = get_name('y', state_name, i_step=i_step)
outputs[out_state_name][i_step + num_starting_times - 1, :] = \
inputs[y_name][0, :]
if is_starting_method:
outputs[starting_name] += np.einsum(
'ij,j...->i...', starting_coeffs[:, i_step, :],
inputs[y_name])
def compute(self, inputs, outputs):
for parameter_name, parameter in iteritems(
self.options['dynamic_parameters']):
in_name = get_name('in', parameter_name)
out_name = get_name('out', parameter_name)
outputs[out_name] = self.mtx.dot(inputs[in_name])
def compute(self, inputs, outputs):
for state_name, state in iteritems(self.options['states']):
initial_condition_name = get_name('initial_condition', state_name)
starting_name = get_name('starting', state_name)
outputs[starting_name] = 0.
outputs[starting_name][0, :] = inputs[initial_condition_name]
def compute_partials(self, inputs, partials):
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
glm_B = self.options['glm_B']
glm_V = self.options['glm_V']
i_step = self.options['i_step']
ii_step = 0
dy_dF = self.dy_dF
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
y_new_name = get_name('y_new', state_name, i_step=i_step)
partials[y_new_name, 'h'][:, 0] = 0.
for j_stage in range(num_stages):
F_name = get_name('F',
state_name,
i_step=i_step,
j_stage=j_stage)
partials[y_new_name,
F_name] = inputs['h'] * dy_dF[state_name, j_stage]
# partials[y_new_name, 'h'][:, 0] += glm_B[ii_step, j_stage] \
# * inputs[F_name].flatten()
partials[y_new_name,
'h'][:,
0] += np.einsum('i,...->i...', glm_B[:, j_stage],
inputs[F_name][0, :]).flatten()
def compute(self, inputs, outputs):
for parameter_name, parameter in iteritems(
self.options['static_parameters']):
in_name = get_name('in', parameter_name)
out_name = get_name('out', parameter_name)
outputs[out_name] = inputs[in_name]
def solve_nonlinear(self, inputs, outputs):
num_times = self.options['num_times']
num_step_vars = self.options['num_step_vars']
glm_B = self.options['glm_B']
dy_dy_inv = self.dy_dy_inv
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name)
y0_name = get_name('y0', state_name)
y_name = get_name('y', state_name)
vec = np.zeros((
num_times,
num_step_vars,
) + shape)
vec[0, :, :] += inputs[y0_name] # y0 term
vec[1:, :, :] += np.einsum('jl,i,il...->ij...', glm_B,
inputs['h_vec'],
inputs[F_name]) # hF term
outputs[y_name] = dy_dy_inv[state_name].solve(vec.flatten(),
'N').reshape((
num_times,
num_step_vars,
) + shape)
def setup(self):
num_step_vars = self.options['num_step_vars']
self.declare_partials('*', '*', dependent=False)
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
initial_condition_name = get_name('initial_condition', state_name)
starting_name = get_name('starting', state_name)
self.add_input(initial_condition_name,
shape=state['shape'],
units=state['units'])
self.add_output(starting_name,
shape=(num_step_vars, ) + state['shape'],
units=state['units'])
ones = np.ones(size)
arange = np.arange(size)
self.declare_partials(starting_name,
initial_condition_name,
val=ones,
rows=arange,
cols=arange)
def apply_nonlinear(self, inputs, outputs, residuals):
num_times = self.options['num_times']
num_step_vars = self.options['num_step_vars']
glm_B = self.options['glm_B']
dy_dy = self.dy_dy
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name)
y0_name = get_name('y0', state_name)
y_name = get_name('y', state_name)
# dy_dy term
in_vec = outputs[y_name].reshape(
(num_times * num_step_vars * size))
out_vec = dy_dy[state_name].dot(in_vec).reshape((
num_times,
num_step_vars,
) + shape)
residuals[y_name] = out_vec # y term
residuals[y_name][0, :, :] -= inputs[y0_name] # y0 term
residuals[y_name][1:, :, :] -= np.einsum('jl,i,il...->ij...',
glm_B, inputs['h_vec'],
inputs[F_name]) # hF term
def compute(self, inputs, outputs):
num_starting_times = self.options['num_starting_times']
num_my_times = self.options['num_my_times']
starting_coeffs = self.options['starting_coeffs']
has_starting_method = num_starting_times > 1
is_starting_method = starting_coeffs is not None
for state_name, state in iteritems(self.options['states']):
y_name = get_name('y', state_name)
starting_state_name = get_name('starting_state', state_name)
out_state_name = get_name('state', state_name)
starting_name = get_name('starting', state_name)
outputs[out_state_name][num_starting_times -
1:] = inputs[y_name][:, 0, :]
if has_starting_method:
outputs[out_state_name][:num_starting_times - 1] = \
inputs[starting_state_name][:-1, :]
if is_starting_method:
outputs[starting_name] = np.einsum('ijk,jk...->i...',
starting_coeffs,
inputs[y_name])
def compute(self, inputs, outputs):
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
i_stage = self.options['i_stage']
glm_A = self.options['glm_A']
glm_U = self.options['glm_U']
i_step = self.options['i_step']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
y_old_name = get_name('y_old',
state_name,
i_step=i_step,
i_stage=i_stage)
Y_name = get_name('Y', state_name, i_step=i_step, i_stage=i_stage)
outputs[Y_name][0, :] = np.einsum('i,i...->...', glm_U[i_stage, :],
inputs[y_old_name])
for j_stage in range(i_stage):
F_name = get_name('F',
state_name,
i_step=i_step,
i_stage=i_stage,
j_stage=j_stage)
outputs[Y_name] += inputs['h'] * glm_A[
i_stage, j_stage] * inputs[F_name]
def compute_partials(self, inputs, partials):
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
i_stage = self.options['i_stage']
glm_A = self.options['glm_A']
glm_U = self.options['glm_U']
i_step = self.options['i_step']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
Y_name = get_name('Y', state_name, i_step=i_step, i_stage=i_stage)
partials[Y_name, 'h'][:, 0] = 0.
for j_stage in range(i_stage):
F_name = get_name('F',
state_name,
i_step=i_step,
i_stage=i_stage,
j_stage=j_stage)
partials[Y_name,
F_name] = inputs['h'] * glm_A[i_stage, j_stage]
partials[Y_name,
'h'][:,
0] += glm_A[i_stage,
j_stage] * inputs[F_name].flatten()
def setup(self):
time_units = self.options['time_units']
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
i_step = self.options['i_step']
glm_B = self.options['glm_B']
glm_V = self.options['glm_V']
self.dy_dF = dy_dF = {}
self.add_input('h', units=time_units)
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name, i_step=i_step)
y_old_name = get_name('y_old', state_name, i_step=i_step)
y_new_name = get_name('y_new', state_name, i_step=i_step)
self.add_input(F_name, shape=(num_stages,) + shape,
units=get_rate_units(state['units'], time_units))
self.add_input(y_old_name, shape=(num_step_vars,) + shape,
units=state['units'])
self.add_output(y_new_name, shape=(num_step_vars,) + shape,
units=state['units'])
F_arange = np.arange(num_stages * size).reshape(
(num_stages,) + shape)
y_arange = np.arange(num_step_vars * size).reshape(
(num_step_vars,) + shape)
# -----------------
# (num_step_vars, num_stages,) + shape
rows = np.einsum('i...,j->ij...', y_arange, np.ones(num_stages, int)).flatten()
cols = np.zeros((num_step_vars, num_stages,) + shape, int).flatten()
self.declare_partials(y_new_name, 'h', rows=rows, cols=cols)
cols = np.einsum('j...,i->ij...', F_arange, np.ones(num_step_vars, int)).flatten()
self.declare_partials(y_new_name, F_name, rows=rows, cols=cols)
# -----------------
# (num_step_vars, num_step_vars,) + shape
data = np.einsum('ij,...->ij...',
glm_V, np.ones(shape)).flatten()
rows = np.einsum('i...,j->ij...', y_arange, np.ones(num_step_vars)).flatten()
cols = np.einsum('j...,i->ij...', y_arange, np.ones(num_step_vars)).flatten()
self.declare_partials(y_new_name, y_old_name, val=data, rows=rows, cols=cols)
def compute(self, inputs, outputs):
glm_A = self.options['glm_A']
glm_U = self.options['glm_U']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name)
y_name = get_name('y', state_name)
Y_out_name = get_name('Y_out', state_name)
Y_in_name = get_name('Y_in', state_name)
outputs[Y_out_name] = -inputs[Y_in_name] \
+ np.einsum('jk,i,ik...->ij...', glm_A, inputs['h_vec'], inputs[F_name]) \
+ np.einsum('jk,ik...->ij...', glm_U, inputs[y_name][:-1, :, :])
def compute(self, inputs, outputs):
normalized_times = self.options['normalized_times']
stage_norm_times = self.options['stage_norm_times']
num_times = len(normalized_times)
num_stage_times = len(stage_norm_times)
for parameter_name, parameter in iteritems(
self.options['dynamic_parameters']):
size = np.prod(parameter['shape'])
shape = parameter['shape']
in_name = get_name('in', parameter_name)
out_name = get_name('out', parameter_name)
outputs[out_name] = self.mtx.dot(inputs[in_name].reshape(
(num_times, size))).reshape((num_stage_times, ) + shape)
def solve_multi_linear(self, d_outputs, d_residuals, mode):
num_times = self.options['num_times']
num_step_vars = self.options['num_step_vars']
dy_dy_inv = self.dy_dy_inv
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
y_name = get_name('y', state_name)
nrow = num_times * num_step_vars * size
ncol = d_outputs[y_name].shape[-1]
if mode == 'fwd':
rhs_array = d_residuals[y_name].reshape((nrow, ncol))
sol_array = d_outputs[y_name].reshape((nrow, ncol))
solve_mode = 'N'
elif mode == 'rev':
rhs_array = d_outputs[y_name].reshape((nrow, ncol))
sol_array = d_residuals[y_name].reshape((nrow, ncol))
solve_mode = 'T'
for icol in range(ncol):
rhs = rhs_array[:, icol]
sol = dy_dy_inv[state_name].solve(rhs, solve_mode)
sol_array[:, icol] = sol
def solve_linear(self, d_outputs, d_residuals, mode):
num_times = self.options['num_times']
num_step_vars = self.options['num_step_vars']
dy_dy_inv = self.dy_dy_inv
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
y_name = get_name('y', state_name)
if mode == 'fwd':
rhs_vec = d_residuals[y_name].flatten()
solve_mode = 'N'
elif mode == 'rev':
rhs_vec = d_outputs[y_name].flatten()
solve_mode = 'T'
sol_vec = dy_dy_inv[state_name].solve(rhs_vec, solve_mode)
if mode == 'fwd':
d_outputs[y_name] = sol_vec.reshape((
num_times,
num_step_vars,
) + shape)
elif mode == 'rev':
d_residuals[y_name] = sol_vec.reshape((
num_times,
num_step_vars,
) + shape)
def compute(self, inputs, outputs):
glm_A = self.options['glm_A']
glm_U = self.options['glm_U']
i_step = self.options['i_step']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name, i_step=i_step)
y_old_name = get_name('y_old', state_name, i_step=i_step)
Y_name = get_name('Y', state_name, i_step=i_step)
outputs[Y_name] = 0. \
+ np.einsum('ij,j...->i...', glm_A, inputs[F_name]) * inputs['h'] \
+ np.einsum('ij,j...->i...', glm_U, inputs[y_old_name])
def setup(self):
normalized_times = self.options['normalized_times']
stage_norm_times = self.options['stage_norm_times']
num_times = len(normalized_times)
num_stage_times = len(stage_norm_times)
data0, rows0, cols0 = get_sparse_linear_spline(normalized_times,
stage_norm_times)
nnz = len(data0)
self.mtx = scipy.sparse.csc_matrix((data0, (rows0, cols0)),
shape=(num_stage_times, num_times))
for parameter_name, parameter in iteritems(
self.options['dynamic_parameters']):
size = np.prod(parameter['shape'])
shape = parameter['shape']
in_name = get_name('in', parameter_name)
out_name = get_name('out', parameter_name)
self.add_input(in_name,
shape=(num_times, ) + shape,
units=parameter['units'])
self.add_output(out_name,
shape=(num_stage_times, ) + shape,
units=parameter['units'])
# (num_stage_times, num_out,) + shape
data = np.einsum('i,...->i...', data0, np.ones(shape)).flatten()
rows = (
np.einsum('i,...->i...', rows0, size * np.ones(shape, int)) +
np.einsum('i,...->i...', np.ones(nnz, int),
np.arange(size).reshape(shape))).flatten()
cols = (
np.einsum('i,...->i...', cols0, size * np.ones(shape, int)) +
np.einsum('i,...->i...', np.ones(nnz, int),
np.arange(size).reshape(shape))).flatten()
self.declare_partials(out_name,
in_name,
val=data,
rows=rows,
cols=cols)
def linearize(self, inputs, outputs, partials):
glm_B = self.options['glm_B']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name)
y0_name = get_name('y0', state_name)
y_name = get_name('y', state_name)
# (num_times - 1, num_step_vars, num_stages,) + shape
partials[y_name,
F_name] = -np.einsum('...,jk,i->ijk...', np.ones(shape),
glm_B, inputs['h_vec']).flatten()
partials[y_name, 'h_vec'] = -np.einsum('jk,ik...->ijk...', glm_B,
inputs[F_name]).flatten()
def setup(self):
for parameter_name, parameter in iteritems(
self.options['static_parameters']):
size = np.prod(parameter['shape'])
shape = parameter['shape']
in_name = get_name('in', parameter_name)
out_name = get_name('out', parameter_name)
self.add_input(in_name, shape=shape, units=parameter['units'])
self.add_output(out_name, shape=shape, units=parameter['units'])
ones = np.ones(size)
arange = np.arange(size)
self.declare_partials(out_name,
in_name,
val=ones,
rows=arange,
cols=arange)
def compute_partials(self, inputs, partials):
num_times = self.options['num_times']
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name)
Y_out_name = get_name('Y_out', state_name)
Y_in_name = get_name('Y_in', state_name)
y0_name = get_name('y0', state_name)
mtx_y0 = self.mtx_y0_dict[state_name]
mtx = self.mtx_dict[state_name]
mtx_h = self.mtx_h_dict[state_name]
partials[Y_out_name, 'h_vec'][:, :] = mtx.dot(np.diag(inputs[F_name].flatten())).dot(mtx_h)
partials[Y_out_name, F_name][:, :] = mtx.dot(np.diag(mtx_h.dot(inputs['h_vec'])))
partials[Y_out_name, y0_name][:, :] = mtx_y0
def compute_partials(self, inputs, partials):
time_units = self.options['time_units']
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
glm_A = self.options['glm_A']
glm_U = self.options['glm_U']
i_step = self.options['i_step']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name, i_step=i_step)
y_old_name = get_name('y_old', state_name, i_step=i_step)
Y_name = get_name('Y', state_name, i_step=i_step)
# (num_stages, num_stages,) + shape
partials[Y_name, F_name] = np.einsum(
'...,ij->ij...', np.ones(shape), glm_A).flatten() * inputs['h']
partials[Y_name, 'h'] = np.einsum('ij,j...->ij...', glm_A,
inputs[F_name]).flatten()
def compute(self, inputs, outputs):
num_times = self.options['num_times']
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name)
Y_out_name = get_name('Y_out', state_name)
Y_in_name = get_name('Y_in', state_name)
y0_name = get_name('y0', state_name)
mtx_y0 = self.mtx_y0_dict[state_name]
mtx = self.mtx_dict[state_name]
mtx_h = self.mtx_h_dict[state_name]
Y_shape = outputs[Y_out_name].shape
outputs[Y_out_name] = -inputs[Y_in_name] \
+ mtx.dot(mtx_h.dot(inputs['h_vec']) * inputs[F_name].flatten()).reshape(Y_shape) \
+ mtx_y0.dot(inputs[y0_name].flatten()).reshape(Y_shape)
def compute_partials(self, inputs, partials):
time_units = self.options['time_units']
num_times = self.options['num_times']
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
glm_A = self.options['glm_A']
glm_U = self.options['glm_U']
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
shape = state['shape']
F_name = get_name('F', state_name)
y_name = get_name('y', state_name)
Y_out_name = get_name('Y_out', state_name)
Y_in_name = get_name('Y_in', state_name)
# (num_times - 1, num_stages, num_stages,) + shape
partials[Y_out_name, F_name] = np.einsum(
'...,jk,i->ijk...', np.ones(shape), glm_A, inputs['h_vec']).flatten()
partials[Y_out_name, 'h_vec'] = np.einsum(
'jk,ik...->ijk...', glm_A, inputs[F_name]).flatten()
def setup(self):
super(ExplicitTMIntegrator, self).setup()
ode_function = self.options['ode_function']
method = self.options['method']
starting_coeffs = self.options['starting_coeffs']
has_starting_method = method.starting_method is not None
is_starting_method = starting_coeffs is not None
states = ode_function._states
static_parameters = ode_function._static_parameters
dynamic_parameters = ode_function._dynamic_parameters
time_units = ode_function._time_options['units']
starting_norm_times, my_norm_times = self._get_meta()
glm_A, glm_B, glm_U, glm_V, num_stages, num_step_vars = self._get_method(
)
num_times = len(my_norm_times)
num_stages = method.num_stages
num_step_vars = method.num_values
glm_A = method.A
glm_B = method.B
glm_U = method.U
glm_V = method.V
# ------------------------------------------------------------------------------------
integration_group = Group()
self.add_subsystem('integration_group', integration_group)
for i_step in range(len(my_norm_times) - 1):
step_comp_old_name = 'integration_group.step_comp_%i' % (i_step -
1)
step_comp_new_name = 'integration_group.step_comp_%i' % (i_step)
for i_stage in range(num_stages):
stage_comp_name = 'integration_group.stage_comp_%i_%i' % (
i_step, i_stage)
ode_comp_name = 'integration_group.ode_comp_%i_%i' % (i_step,
i_stage)
if ode_function._time_options['targets']:
self.connect('time_comp.stage_times', [
'.'.join((ode_comp_name, t))
for t in ode_function._time_options['targets']
],
src_indices=i_step * (num_stages) + i_stage)
comp = ExplicitTMStageComp(
states=states,
time_units=time_units,
num_stages=num_stages,
num_step_vars=num_step_vars,
glm_A=glm_A,
glm_U=glm_U,
i_stage=i_stage,
i_step=i_step,
)
integration_group.add_subsystem(
stage_comp_name.split('.')[1], comp)
self.connect('time_comp.h_vec',
'%s.h' % stage_comp_name,
src_indices=i_step)
for j_stage in range(i_stage):
ode_comp_tmp_name = 'integration_group.ode_comp_%i_%i' % (
i_step, j_stage)
self._connect_multiple(
self._get_state_names(ode_comp_tmp_name,
'rate_source'),
self._get_state_names(stage_comp_name,
'F',
i_step=i_step,
i_stage=i_stage,
j_stage=j_stage),
)
comp = self._create_ode(1)
integration_group.add_subsystem(
ode_comp_name.split('.')[1], comp)
self._connect_multiple(
self._get_state_names(stage_comp_name,
'Y',
i_step=i_step,
i_stage=i_stage),
self._get_state_names(ode_comp_name, 'targets'),
)
if len(static_parameters) > 0:
self._connect_multiple(
self._get_static_parameter_names(
'static_parameter_comp', 'out'),
self._get_static_parameter_names(
ode_comp_name, 'targets'),
)
if len(dynamic_parameters) > 0:
src_indices_list = []
for parameter_name, value in iteritems(dynamic_parameters):
size = np.prod(value['shape'])
shape = value['shape']
arange = np.arange(((len(my_norm_times) - 1) *
num_stages * size)).reshape(((
len(my_norm_times) - 1,
num_stages,
) + shape))
src_indices = arange[i_step,
i_stage, :].reshape((1, ) + shape)
src_indices_list.append(src_indices)
self._connect_multiple(
self._get_dynamic_parameter_names(
'dynamic_parameter_comp', 'out'),
self._get_dynamic_parameter_names(
ode_comp_name, 'targets'),
src_indices_list,
)
comp = ExplicitTMStepComp(
states=states,
time_units=time_units,
num_stages=num_stages,
num_step_vars=num_step_vars,
glm_B=glm_B,
glm_V=glm_V,
i_step=i_step,
)
integration_group.add_subsystem(
step_comp_new_name.split('.')[1], comp)
self.connect('time_comp.h_vec',
'%s.h' % step_comp_new_name,
src_indices=i_step)
for j_stage in range(num_stages):
ode_comp_tmp_name = 'integration_group.ode_comp_%i_%i' % (
i_step, j_stage)
self._connect_multiple(
self._get_state_names(ode_comp_tmp_name, 'rate_source'),
self._get_state_names(step_comp_new_name,
'F',
i_step=i_step,
j_stage=j_stage),
)
if i_step == 0:
self._connect_multiple(
self._get_state_names('starting_system', 'starting'),
self._get_state_names(step_comp_new_name,
'y_old',
i_step=i_step),
)
for i_stage in range(num_stages):
stage_comp_name = 'integration_group.stage_comp_%i_%i' % (
i_step, i_stage)
self._connect_multiple(
self._get_state_names('starting_system', 'starting'),
self._get_state_names(stage_comp_name,
'y_old',
i_step=i_step,
i_stage=i_stage),
)
else:
self._connect_multiple(
self._get_state_names(step_comp_old_name,
'y_new',
i_step=i_step - 1),
self._get_state_names(step_comp_new_name,
'y_old',
i_step=i_step),
)
for i_stage in range(num_stages):
stage_comp_name = 'integration_group.stage_comp_%i_%i' % (
i_step, i_stage)
self._connect_multiple(
self._get_state_names(step_comp_old_name,
'y_new',
i_step=i_step - 1),
self._get_state_names(stage_comp_name,
'y_old',
i_step=i_step,
i_stage=i_stage),
)
promotes_outputs = []
for state_name in states:
out_state_name = get_name('state', state_name)
starting_name = get_name('starting', state_name)
promotes_outputs.append(out_state_name)
if is_starting_method:
promotes_outputs.append(starting_name)
comp = TMOutputComp(states=states,
num_starting_times=len(starting_norm_times),
num_my_times=len(my_norm_times),
num_step_vars=num_step_vars,
starting_coeffs=starting_coeffs)
self.add_subsystem('output_comp',
comp,
promotes_outputs=promotes_outputs)
if has_starting_method:
self._connect_multiple(
self._get_state_names('starting_system', 'state'),
self._get_state_names('output_comp', 'starting_state'),
)
for i_step in range(len(my_norm_times)):
if i_step == 0:
self._connect_multiple(
self._get_state_names('starting_system', 'starting'),
self._get_state_names('output_comp', 'y', i_step=i_step),
)
else:
self._connect_multiple(
self._get_state_names('integration_group.step_comp_%i' %
(i_step - 1),
'y_new',
i_step=i_step - 1),
self._get_state_names('output_comp', 'y', i_step=i_step),
)
def setup(self):
super(ImplicitTMIntegrator, self).setup()
ode_function = self.options['ode_function']
method = self.options['method']
starting_coeffs = self.options['starting_coeffs']
has_starting_method = method.starting_method is not None
is_starting_method = starting_coeffs is not None
states = ode_function._states
static_parameters = ode_function._static_parameters
dynamic_parameters = ode_function._dynamic_parameters
time_units = ode_function._time_options['units']
starting_norm_times, my_norm_times = self._get_meta()
glm_A, glm_B, glm_U, glm_V, num_stages, num_step_vars = self._get_method()
num_times = len(my_norm_times)
num_stages = method.num_stages
num_step_vars = method.num_values
glm_A = method.A
glm_B = method.B
glm_U = method.U
glm_V = method.V
# ------------------------------------------------------------------------------------
integration_group = Group()
self.add_subsystem('integration_group', integration_group)
for i_step in range(len(my_norm_times) - 1):
group = Group(assembled_jac_type='dense')
group_old_name = 'integration_group.step_%i' % (i_step - 1)
group_new_name = 'integration_group.step_%i' % i_step
integration_group.add_subsystem(group_new_name.split('.')[1], group)
comp = self._create_ode(num_stages)
group.add_subsystem('ode_comp', comp)
if ode_function._time_options['targets']:
self.connect('time_comp.stage_times',
['.'.join((group_new_name + '.ode_comp', t)) for t in
ode_function._time_options['targets']],
src_indices=i_step * (num_stages) + np.arange(num_stages))
if len(static_parameters) > 0:
self._connect_multiple(
self._get_static_parameter_names('static_parameter_comp', 'out'),
self._get_static_parameter_names(group_new_name + '.ode_comp', 'targets'),
src_indices_list=[[0] * num_stages for _ in range(len(static_parameters))]
)
if len(dynamic_parameters) > 0:
src_indices_list = []
for parameter_name, value in iteritems(dynamic_parameters):
size = np.prod(value['shape'])
shape = value['shape']
arange = np.arange(((len(my_norm_times) - 1) * num_stages * size)).reshape(
((len(my_norm_times) - 1, num_stages,) + shape))
src_indices = arange[i_step, :, :]
src_indices_list.append(src_indices.flat)
self._connect_multiple(
self._get_dynamic_parameter_names('dynamic_parameter_comp', 'out'),
self._get_dynamic_parameter_names(group_new_name + '.ode_comp', 'targets'),
src_indices_list,
)
comp = ImplicitTMStageComp(
states=states, time_units=time_units,
num_stages=num_stages, num_step_vars=num_step_vars,
glm_A=glm_A, glm_U=glm_U, i_step=i_step,
)
group.add_subsystem('stage_comp', comp)
self.connect('time_comp.h_vec', group_new_name + '.stage_comp.h', src_indices=i_step)
comp = ImplicitTMStepComp(
states=states, time_units=time_units,
num_stages=num_stages, num_step_vars=num_step_vars,
glm_B=glm_B, glm_V=glm_V, i_step=i_step,
)
group.add_subsystem('step_comp', comp)
self.connect('time_comp.h_vec', group_new_name + '.step_comp.h', src_indices=i_step)
self._connect_multiple(
self._get_state_names(group_new_name + '.ode_comp', 'rate_source'),
self._get_state_names(group_new_name + '.step_comp', 'F', i_step=i_step),
)
self._connect_multiple(
self._get_state_names(group_new_name + '.ode_comp', 'rate_source'),
self._get_state_names(group_new_name + '.stage_comp', 'F', i_step=i_step),
)
self._connect_multiple(
self._get_state_names(group_new_name + '.stage_comp', 'Y', i_step=i_step),
self._get_state_names(group_new_name + '.ode_comp', 'targets'),
)
if i_step == 0:
self._connect_multiple(
self._get_state_names('starting_system', 'starting'),
self._get_state_names(group_new_name + '.step_comp', 'y_old', i_step=i_step),
)
self._connect_multiple(
self._get_state_names('starting_system', 'starting'),
self._get_state_names(group_new_name + '.stage_comp', 'y_old', i_step=i_step),
)
else:
self._connect_multiple(
self._get_state_names(group_old_name + '.step_comp', 'y_new', i_step=i_step - 1),
self._get_state_names(group_new_name + '.step_comp', 'y_old', i_step=i_step),
)
self._connect_multiple(
self._get_state_names(group_old_name + '.step_comp', 'y_new', i_step=i_step - 1),
self._get_state_names(group_new_name + '.stage_comp', 'y_old', i_step=i_step),
)
group.nonlinear_solver = NewtonSolver(iprint=2, maxiter=100)
group.linear_solver = DirectSolver(assemble_jac=True)
promotes = []
promotes.extend([get_name('state', state_name) for state_name in states])
if is_starting_method:
promotes.extend([get_name('starting', state_name) for state_name in states])
comp = TMOutputComp(
states=states, num_starting_times=len(starting_norm_times),
num_my_times=len(my_norm_times), num_step_vars=num_step_vars,
starting_coeffs=starting_coeffs)
self.add_subsystem('output_comp', comp, promotes_outputs=promotes)
if has_starting_method:
self._connect_multiple(
self._get_state_names('starting_system', 'state'),
self._get_state_names('output_comp', 'starting_state'),
)
for i_step in range(len(my_norm_times)):
if i_step == 0:
self._connect_multiple(
self._get_state_names('starting_system', 'starting'),
self._get_state_names('output_comp', 'y', i_step=i_step),
)
else:
self._connect_multiple(
self._get_state_names('integration_group.step_%i' % (i_step - 1) + '.step_comp', 'y_new', i_step=i_step - 1),
self._get_state_names('output_comp', 'y', i_step=i_step),
)
def setup(self):
num_starting_times = self.options['num_starting_times']
num_my_times = self.options['num_my_times']
num_step_vars = self.options['num_step_vars']
starting_coeffs = self.options['starting_coeffs']
num_times = num_starting_times + num_my_times - 1
has_starting_method = num_starting_times > 1
is_starting_method = starting_coeffs is not None
if is_starting_method:
num_starting = starting_coeffs.shape[0]
self.declare_partials('*', '*', dependent=False)
for state_name, state in iteritems(self.options['states']):
shape = state['shape']
size = np.prod(shape)
starting_state_name = get_name('starting_state', state_name)
out_state_name = get_name('state', state_name)
starting_name = get_name('starting', state_name)
for i_step in range(num_my_times):
y_name = get_name('y', state_name, i_step=i_step)
self.add_input(y_name,
shape=(num_step_vars, ) + shape,
units=state['units'])
if has_starting_method:
self.add_input(starting_state_name,
shape=(num_starting_times, ) + shape,
units=state['units'])
self.add_output(out_state_name,
shape=(num_times, ) + state['shape'],
units=state['units'])
if is_starting_method:
self.add_output(starting_name,
shape=(num_starting, ) + shape,
units=state['units'])
y_arange = np.arange(num_step_vars *
size).reshape((num_step_vars, ) + shape)
state_arange = np.arange(num_times * size).reshape((num_times, ) +
shape)
if has_starting_method:
out_state_arange = np.arange(
num_times * size).reshape((num_times, ) + shape)
if is_starting_method:
starting_arange = np.arange(
num_starting * size).reshape((num_starting, ) + shape)
if has_starting_method:
starting_state_arange = np.arange(
num_starting_times * size).reshape((num_starting_times, ) +
shape)
data = np.ones((num_starting_times - 1) * size, int)
rows = out_state_arange[:num_starting_times - 1, :].flatten()
cols = starting_state_arange[:-1, :].flatten()
self.declare_partials(out_state_name,
starting_state_name,
val=data,
rows=rows,
cols=cols)
for i_step in range(num_my_times):
y_name = get_name('y', state_name, i_step=i_step)
data = np.ones(size)
rows = (state_arange[i_step + num_starting_times -
1, :]).flatten()
cols = (y_arange[0, :]).flatten()
self.declare_partials(out_state_name,
y_name,
val=data,
rows=rows,
cols=cols)
if is_starting_method:
# (num_starting, num_step_vars,) + shape
data = np.einsum('ij,...->ij...',
starting_coeffs[:, i_step, :],
np.ones(shape)).flatten()
rows = np.einsum('j,i...->ij...',
np.ones(num_step_vars, int),
starting_arange).flatten()
cols = np.einsum('i,j...->ij...',
np.ones(num_starting,
int), y_arange).flatten()
self.declare_partials(starting_name,
y_name,
val=data,
rows=rows,
cols=cols)
def setup(self):
time_units = self.options['time_units']
num_stages = self.options['num_stages']
num_step_vars = self.options['num_step_vars']
i_step = self.options['i_step']
glm_B = self.options['glm_B']
glm_V = self.options['glm_V']
self.dy_dF = dy_dF = {}
self.declare_partials('*', '*', dependent=False)
self.add_input('h', units=time_units)
for state_name, state in iteritems(self.options['states']):
size = np.prod(state['shape'])
y_old_name = get_name('y_old', state_name, i_step=i_step)
y_new_name = get_name('y_new', state_name, i_step=i_step)
for j_stage in range(num_stages):
F_name = get_name('F',
state_name,
i_step=i_step,
j_stage=j_stage)
self.add_input(F_name,
shape=(1, ) + state['shape'],
units=get_rate_units(state['units'],
time_units))
self.add_input(y_old_name,
shape=(num_step_vars, ) + state['shape'],
units=state['units'])
self.add_output(y_new_name,
shape=(num_step_vars, ) + state['shape'],
units=state['units'])
self.declare_partials(y_new_name, 'h', dependent=True)
y_arange = np.arange(num_step_vars * size).reshape(
(num_step_vars, size))
# num_step_vars, num_step_vars, size
data = np.einsum('ij,...->ij...', glm_V, np.ones(size)).flatten()
rows = np.einsum('i...,j->ij...', y_arange,
np.ones(num_step_vars, int)).flatten()
cols = np.einsum('j...,i->ij...', y_arange,
np.ones(num_step_vars, int)).flatten()
self.declare_partials(y_new_name,
y_old_name,
val=data,
rows=rows,
cols=cols)
for j_stage in range(num_stages):
F_name = get_name('F',
state_name,
i_step=i_step,
j_stage=j_stage)
vals = np.zeros((num_step_vars, 1, size))
rows = np.arange(num_step_vars * 1 * size)
cols = np.zeros((num_step_vars, 1 * size), int)
for ii_step in range(num_step_vars):
vals[ii_step, 0, :] = glm_B[ii_step, j_stage]
cols[ii_step, :] = np.arange(1 * size)
vals = vals.flatten()
cols = cols.flatten()
self.declare_partials(y_new_name, F_name, rows=rows, cols=cols)
dy_dF[state_name, j_stage] = vals
