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