Changed function format

docs/src/tutorials/quadcopter.md

@@ -42,7 +42,7 @@ using LinearAlgebra
 
 ### 2. Function Design
 
-We will define a function called `Quad` that will take arguments for the number of nodes and the discretization size (i.e., $\Delta t$), optimize the model, and return the graph and reference values $x^{ref}$. 
+We will define a function called `build_quadcopter_graph` that will take arguments for the number of nodes and the discretization size (i.e., $\Delta t$), optimize the model, and return the graph and reference values $x^{ref}$.  
 
 The function inputs are:
 - number of nodes (`N`)
@@ -53,64 +53,28 @@ The function outputs are:
 - The objective value of the discretized form of $\phi$
 - An array with the reference values on each node ($x^{ref}$)
 
-We first establish the set points for each timepoint. The quadcopter will fly in a linear upward path in the positive X, Y, and Z directions.
-```julia
-function Quad(N, dt)
-    X_ref = 0:10/N:10;
-    dXdt_ref = zeros(N);
-    Y_ref = 0:10/N:10;
-    dYdt_ref = zeros(N)
-    Z_ref = 0:10/N:10;
-    dZdt_ref = zeros(N);
-    g_ref = zeros(N);
-    b_ref = zeros(N);
-    a_ref = zeros(N);
-```    
-
-We next combine these vectors into another vector which we call $x^{ref}$, and then define the necessary constants and arrays.
+The `build_quadcopter_graph` function will use three supporting functions that will add variables, add constraints (both local and linking) and add the objectives to the nodes. These functions will be detailed below before they are used to build the full quadcopter graph. 
 
+The first function we call `add_variables!`, which defines each of the decision variables as well as some supporting expressions. Here, we loop through the nodes, and define variables on each node. These variables include expressions that will simplify forming the linking constrints (these are the right hand sides of the derivatives). 
 ```julia
-    xk_ref = [X_ref, dXdt_ref, Y_ref, dYdt_ref, Z_ref, dZdt_ref, g_ref, b_ref, a_ref];
+function add_variables!(nodes)
     grav = 9.8 # m/s^2
-
-    Q = diagm([1, 0, 1, 0, 1, 0, 1, 1, 1]);
-    R = diagm([1/10, 1/10, 1/10, 1/10]);
 
-    xk_ref1 = zeros(N,9)
-    for i in (1:N)
-        for j in 1:length(xk_ref)
-            xk_ref1[i,j] = xk_ref[j][i]
-        end
-    end
-```
-
-We can now begin building the OptiGraph. We now initialize an OptiGraph and set the optimizer. We then define `N` nodes on the OptiGraph `graph`. 
-```Julia
-    graph = OptiGraph()
-    solver = optimizer_with_attributes(Ipopt.Optimizer, "max_iter" => 100)
-    set_optimizer(graph, solver)
-    @optinode(graph, nodes[1:N])
-```
-
-Next, we loop through the nodes, and define variables on each node. These variables include dummy variables for the second derivatives and the angles' derivatives to avoid nonlinearities in the linking constraints. We also define the objective function and define arrays for the variables (to simplify the objective function formulation). 
-```julia        
     for (i, node) in enumerate(nodes)
-
         # Create state variables
         @variable(node, g)
         @variable(node, b)
         @variable(node, a)
-    
+
         @variable(node, X)
         @variable(node, Y)
         @variable(node, Z)
-    
+
         @variable(node, dXdt)
         @variable(node, dYdt)
         @variable(node, dZdt)
 
         # Create input variables
-
         @variable(node, C_a)
         @variable(node, wx)
         @variable(node, wy)
@@ -124,17 +88,16 @@ Next, we loop through the nodes, and define variables on each node. These variab
         @expression(node, dgdt, (wx * cos(g) + wy * sin(g)) / (cos(b)))
         @expression(node, dbdt, - wx * sin(g) + wy * cos(g))
         @expression(node, dadt, wx * cos(g) * tan(b) + wy * sin(g) * tan(b) + wz)
-
-        xk = [X, dXdt, Y, dYdt, Z, dZdt, g, b, a] # Array to hold variables
-        xk1 = xk .- xk_ref1[i,:] # Array to hold the difference between variable values and their setpoints.
-
-        uk = [C_a, wx, wy, wz]
-        @objective(node, Min, (1 / 2 * (xk1') * Q * (xk1) + 1 / 2 * (uk') * R * (uk)) * dt)
     end
-```
-We also set initial conditions on the first variables (for the differential equations). Note that the variables can be referenced by calling `nodes[1][<variable_reference>]`
+end
+```    
+
+Next, we define a function for adding the constraints to the graph. We will set the initial values at time 0 and then define the linking constraints, which are discretized derivatives. Note that both linear and nonlinear constraints are handled in the same way by the user in both the `@constraint` and `@linkconstraint` macros. 
 
 ```julia
+function add_constraints!(graph, nodes, dt)
+    N = length(nodes)
+
     @constraint(nodes[1], nodes[1][:X] == 0)
     @constraint(nodes[1], nodes[1][:Y] == 0)
     @constraint(nodes[1], nodes[1][:Z] == 0)
@@ -144,36 +107,91 @@ We also set initial conditions on the first variables (for the differential equa
     @constraint(nodes[1], nodes[1][:g] == 0)
     @constraint(nodes[1], nodes[1][:b] == 0)
     @constraint(nodes[1], nodes[1][:a] == 0)
-```            
 
-With the variables defined, we can now add linking constraints between each time point. Note that expressions like `d2Xdt2` are the right hand side of the constraints for the derivative of `dXdt`.
-
-```julia
     for i in 1:(N-1) # iterate through each node except the last
         @linkconstraint(graph, nodes[i+1][:dXdt] == dt*nodes[i][:d2Xdt2] + nodes[i][:dXdt])
         @linkconstraint(graph, nodes[i+1][:dYdt] == dt*nodes[i][:d2Ydt2] + nodes[i][:dYdt])
         @linkconstraint(graph, nodes[i+1][:dZdt] == dt*nodes[i][:d2Zdt2] + nodes[i][:dZdt])
-    
+
         @linkconstraint(graph, nodes[i+1][:g] == dt*nodes[i][:dgdt] + nodes[i][:g])
         @linkconstraint(graph, nodes[i+1][:b] == dt*nodes[i][:dbdt] + nodes[i][:b])
         @linkconstraint(graph, nodes[i+1][:a] == dt*nodes[i][:dadt] + nodes[i][:a])
-    
+
         @linkconstraint(graph, nodes[i+1][:X] == dt*nodes[i][:dXdt] + nodes[i][:X])
         @linkconstraint(graph, nodes[i+1][:Y] == dt*nodes[i][:dYdt] + nodes[i][:Y])
         @linkconstraint(graph, nodes[i+1][:Z] == dt*nodes[i][:dZdt] + nodes[i][:Z])
     end
+end
 ```
 
-We now have all the constraints and variables defined and can call `optimize!` on the graph. 
+Next, we set a function for defining the objectives. The quadcopter will fly in a linear upward path in the positive X, Y, and Z directions. We combine these vectors into another vector which we call $x^{ref}$, and then define the necessary constants and arrays (this will simplify forming the objective function).
 
 ```julia
+function add_objective!(nodes, N, dt)
+    X_ref = 0:10/N:10;
+    dXdt_ref = zeros(N);
+    Y_ref = 0:10/N:10;
+    dYdt_ref = zeros(N)
+    Z_ref = 0:10/N:10;
+    dZdt_ref = zeros(N);
+    g_ref = zeros(N);
+    b_ref = zeros(N);
+    a_ref = zeros(N);
+
+    xk_ref = [X_ref, dXdt_ref, Y_ref, dYdt_ref, Z_ref, dZdt_ref, g_ref, b_ref, a_ref];
+
+    Q = diagm([1, 0, 1, 0, 1, 0, 1, 1, 1]);
+    R = diagm([1/10, 1/10, 1/10, 1/10]);
+
+    xk_ref1 = zeros(N,9)
+    for i in (1:N)
+        for j in 1:length(xk_ref)
+            xk_ref1[i,j] = xk_ref[j][i]
+        end
+    end
+
+    for (i, node) in enumerate(nodes)
+        xk = [ # Array to hold variables
+            node[:X],
+            node[:dXdt],
+            node[:Y],
+            node[:dYdt],
+            node[:Z],
+            node[:dZdt],
+            node[:g],
+            node[:b],
+            node[:a]
+        ]
+
+        xk1 = xk .- xk_ref1[i, :] # Array to hold the difference between variable values and their setpoints.
+
+        uk = [node[:C_a], node[:wx], node[:wy], node[:wz]]
+        @objective(node, Min, (1 / 2 * (xk1') * Q * (xk1) + 1 / 2 * (uk') * R * (uk)) * dt)
+    end
+    return xk_ref
+end
+```
+
+We can now define the function for building the optigraph. We initialize an OptiGraph and set the optimizer. We then define `N` nodes on the OptiGraph `graph`. We then call the three functions above and call `set_to_node_objectives` to set the graph's objective to the nodes' objectives we have defined. We can then call `optimize` and return the objective value, the graph, and the reference points. 
+```Julia
+function build_quadcopter_graph(N, dt)
+    graph = OptiGraph()
+    solver = optimizer_with_attributes(Ipopt.Optimizer, "max_iter" => 100)
+    set_optimizer(graph, solver)
+    @optinode(graph, nodes[1:N])
+
+    add_variables!(nodes)
+    add_constraints!(graph, nodes, dt)
+    xk_ref = add_objective!(nodes, N, dt)
+
     set_to_node_objectives(graph)
-    
+
     optimize!(graph);
 
     return objective_value(graph), graph, xk_ref
 end
 ```
+
 Now that we have created our function to model the behavior of the quadcopter,
 we can test it using some example cases.
 
@@ -183,7 +201,7 @@ First, we will run an example with 50 time points (each represented by a node) w
 ```julia
 N = 50
 dt = 0.1
-objv, graph, xk_ref = Quad(N, dt)
+objv, graph, xk_ref = build_quadcopter_graph(N, dt)
 nodes = getnodes(graph)
 # create empty arrays
 CAval_array = zeros(length(nodes))
@@ -237,7 +255,7 @@ obj_val_N = zeros(N)
 
 for i in 1:length(time_steps)
     timing = @elapsed begin
-    objval, graph, xk_ref = Quad(time_steps[i], 10 / time_steps[i]);
+    objval, graph, xk_ref = build_quadcopter_graph(time_steps[i], 10 / time_steps[i]);
     obj_val_N[i] = objval
     end
     println("Done with iteration $i after ", timing, " seconds")