diff --git a/opty/direct_collocation.py b/opty/direct_collocation.py index 4c084b23..50ef2be3 100644 --- a/opty/direct_collocation.py +++ b/opty/direct_collocation.py @@ -597,6 +597,43 @@ def plot_objective_value(self): return ax + def parse_free(self, free): + """Parses the free parameters vector and returns it's components. + + Parameters + ---------- + free : ndarray, shape(n*N + q*N + r + s) + The free parameters of the system. + + where: + + - N : number of collocation nodes + - n : number of unknown state trajectories + - q : number of unknown input trajectories + - r : number of unknown parameters + - s : number of unknown time intervals (s=1 if variable duration, else s=0) + + + Returns + ------- + states : ndarray, shape(n, N) + The array of n states through N time steps. + specified_values : ndarray, shape(q, N) or shape(N,), or None + The array of q specified inputs through N time steps. + constant_values : ndarray, shape(r,) + The array of r constants. + time_interval : float + The time between collocation nodes. Only returned if + ``variable_duration`` is ``True``. + + """ + + n = self.collocator.num_states + N = self.collocator.num_collocation_nodes + q = self.collocator.num_unknown_input_trajectories + variable_duration = self.collocator._variable_duration + + return parse_free(free, n, q, N, variable_duration) class ConstraintCollocator(object): """This class is responsible for generating the constraint function and the diff --git a/opty/tests/test_direct_collocation.py b/opty/tests/test_direct_collocation.py index adc2f3f8..ee598876 100644 --- a/opty/tests/test_direct_collocation.py +++ b/opty/tests/test_direct_collocation.py @@ -10,7 +10,7 @@ from pytest import raises from ..direct_collocation import Problem, ConstraintCollocator -from ..utils import create_objective_function, sort_sympy +from ..utils import create_objective_function, sort_sympy, parse_free def test_pendulum(): @@ -1566,3 +1566,124 @@ def test_for_algebraic_eoms(): ) assert excinfo.type is ValueError + + +def test_prob_parse_free(): + """ + Test for parse_free method of Problem class. + =========================================== + + This test whether the parse_free method of the Problem class works as + the parse_free in utils. + + **States** + + - :math:`x_1, x_2, ux_1, ux_2` : state variables + + **Control** + + - :math:`u_1, u_2` : control variable + + """ + + t = mech.dynamicsymbols._t + + x1, x2, ux1, ux2 = mech.dynamicsymbols('x1 x2 ux1 ux2') + u1, u2 = mech.dynamicsymbols('u1 u2') + h = sym.symbols('h') + a, b = sym.symbols('a b') + + # equations of motion. + # (No meaning, just for testing) + eom = sym.Matrix([ + -x1.diff(t) + ux1, + -x2.diff(t) + ux2, + -ux1.diff(t) + a*u1, + -ux2.diff(t) + b*u2, + ]) + + # Set up and Solve the Optimization Problem + num_nodes = 11 + t0, tf = 0.0, 0.9 + state_symbols = (x1, x2, ux1, ux2) + control_symbols = (u1, u2) + + interval_value = (tf - t0)/(num_nodes - 1) + times = np.linspace(t0, tf, num_nodes) + + # Specify the symbolic instance constraints, as per the example. + instance_constraints = ( + x1.func(t0) - 1.0, + x2.func(t0) + 1.0, + ) + + # Specify the objective function and form the gradient. + + def obj(free): + return sum([free[i]**2 for i in range(2*num_nodes)]) + + def obj_grad(free): + grad = np.zeros_like(free) + grad[:2*num_nodes] = 2*free[:2*num_nodes] + return grad + + # Create the optimization problem and set any options. + prob = Problem( + obj, + obj_grad, + eom, + state_symbols, + num_nodes, + interval_value, + instance_constraints=instance_constraints, +) + + # Give some estimates for the trajectory. + initial_guess = np.random.rand(prob.num_free) + initial_guess1 = initial_guess + + # check whether same results. + statesu, controlsu, constantsu = parse_free(initial_guess1, + len(state_symbols), len(control_symbols), num_nodes) + + states, controls, constants = prob.parse_free(initial_guess) + np.testing.assert_allclose(states, statesu) + np.testing.assert_allclose(controls, controlsu) + np.testing.assert_allclose(constants, constantsu) + + # test with variable interval_value + interval_value = h + t0, tf = 0.0, (num_nodes - 1)*interval_value + def obj(free): + return sum([free[i]**2 for i in range(2*num_nodes)]) + + def obj_grad(free): + grad = np.zeros_like(free) + grad[:2*num_nodes] = 2*free[:2*num_nodes] + return grad + + # Create the optimization problem and set any options. + prob = Problem( + obj, + obj_grad, + eom, + state_symbols, + num_nodes, + interval_value, + instance_constraints=instance_constraints, +) + + # Give some estimates for the trajectory. + initial_guess = np.random.rand(prob.num_free) + initial_guess1 = initial_guess + + # check whether same results. + statesu, controlsu, constantsu, timeu = parse_free(initial_guess1, + len(state_symbols), len(control_symbols), + num_nodes, variable_duration=True) + + states, controls, constants, times = prob.parse_free(initial_guess) + np.testing.assert_allclose(states, statesu) + np.testing.assert_allclose(controls, controlsu) + np.testing.assert_allclose(constants, constantsu) + np.testing.assert_allclose(timeu, times)