|
| 1 | +```@meta |
| 2 | +CurrentModule = OpenMDAODocs |
| 3 | +``` |
| 4 | +# Variable Shapes at Runtime: `shape_by_conn` and `copy_shape` |
| 5 | + |
| 6 | +## A Simple Example |
| 7 | +OpenMDAO is able to [determine variable shapes at runtime](https://openmdao.org/newdocs/versions/latest/features/experimental/dyn_shapes.html). |
| 8 | +In "normal" (aka non-Julian) OpenMDAO, this is done via the `shape_by_conn` and `copy_shape` arguments to the venerable `add_input` and/or `add_output` `Component` methods. |
| 9 | +In OpenMDAO.jl, we can provide the `shape_by_conn` and/or `copy_shape` arguments to the `VarData` `struct` constructor to get the same behavior. |
| 10 | + |
| 11 | +We'll show how this works using a simple `ExplicitComponent` that computes ``y = 2*x^2 + 1`` element-wise, where ``x`` and ``y`` are two-dimensional arrays of any (identical) size. |
| 12 | + |
| 13 | +We'll need `OpenMDAOCore` of course, and need to declare our `ExplicitComponent` in the usual way: |
| 14 | + |
| 15 | +```@example shape_by_conn1 |
| 16 | +using OpenMDAOCore: OpenMDAOCore |
| 17 | +
|
| 18 | +struct ECompShapeByConn <: OpenMDAOCore.AbstractExplicitComp end |
| 19 | +``` |
| 20 | + |
| 21 | +Next we need a `setup` method: |
| 22 | + |
| 23 | +```@example shape_by_conn1 |
| 24 | +function OpenMDAOCore.setup(self::ECompShapeByConn) |
| 25 | + input_data = [OpenMDAOCore.VarData("x"; shape_by_conn=true)] |
| 26 | + output_data = [OpenMDAOCore.VarData("y"; shape_by_conn=true, copy_shape="x")] |
| 27 | +
|
| 28 | + partials_data = [] |
| 29 | + return input_data, output_data, partials_data |
| 30 | +end |
| 31 | +``` |
| 32 | + |
| 33 | +Notice how we provided the `shape_by_conn` argument to the `VarData` `struct` for `x`, and the `shape_by_conn` and `copy_shape` arguments to `y`'s `VarData` `struct`. |
| 34 | +This means that the shape of `x` will be determined at runtime by OpenMDAO, and will be set to the shape of whatever output is connected to `x`. |
| 35 | +The shape of `y` will be set to that of `x`, since we provided the `copy_shape="x"` argument. |
| 36 | +(Also notice how we returned an empty Vector for the `partials_data` output—OpenMDAO.jl always expects `OpenMDAOCore.setup` to return three Vectors, corresponding to `input_data`, `output_data`, and `partials_data`. |
| 37 | +But the `partials_data` Vector can be empty if it's not needed.) |
| 38 | + |
| 39 | +Now, the derivative of `y` with respect to `x` will be sparse—the value of an element `y[i,j]` depends on the element `x[i,j]`, and no others. |
| 40 | +We can communicate this fact to OpenMDAO through the `rows` and `cols` arguments to `declare_partials` in Python OpenMDAO, or the `PartialsData` `struct` in OpenMDAO.jl. |
| 41 | +But how do we do that here, when we don't know the sizes of `x` and `y` in the `setup` method? |
| 42 | +The answer is we implement an `OpenMDAOCore.setup_partials` method, which gives us another chance to create more `PartialsData` `structs` after OpenMDAO has figured out what the sizes of all the inputs and outputs are: |
| 43 | + |
| 44 | +```@example shape_by_conn1 |
| 45 | +function OpenMDAOCore.setup_partials(self::ECompShapeByConn, input_sizes, output_sizes) |
| 46 | + @assert input_sizes["x"] == output_sizes["y"] |
| 47 | + m, n = input_sizes["x"] |
| 48 | + rows, cols = OpenMDAOCore.get_rows_cols(ss_sizes=Dict(:i=>m, :j=>n), of_ss=[:i, :j], wrt_ss=[:i, :j]) |
| 49 | + partials_data = [OpenMDAOCore.PartialsData("y", "x"; rows=rows, cols=cols)] |
| 50 | +
|
| 51 | + return partials_data |
| 52 | +end |
| 53 | +``` |
| 54 | + |
| 55 | +The `OpenMDAOCore.setup_partials` method will always take an instance of the `OpenMDAOCore.AbstractComp` (called `self` here), and two `Dict`s, both with `String` keys and `NTuple{N, Int}` values. |
| 56 | +The keys indicate the name of an input or output variable, and the `NTuple{Int, N}` values are the shapes of each variable. |
| 57 | +The first `Dict` holds all the input shapes, and the second `Dict` has all the output shapes. |
| 58 | + |
| 59 | +Now, the job of `setup_partials` is to return a `Vector` of `PartialsData` `structs`. |
| 60 | +We'd like to include the `rows` and `cols` arguments to the `PartialsData` `struct` for the derivative of `y` with respect to `x`, but it's a bit tricky, since `x` and `y` are two-dimensional. |
| 61 | +Luckily, there is a small utility function provided by OpenMDAOCore.jl called `get_rows_cols` that can help us. |
| 62 | + |
| 63 | +## Sparsity Patterns with `get_rows_cols` |
| 64 | +The `get_rows_cols` function uses a symbolic notation to express sparsity patterns in a simple way. |
| 65 | +Here's an example that corresponds to our present case. |
| 66 | +Let's say `x` and `y` have shape `(2, 3)`. |
| 67 | +Then the non-zero index combinations for the derivative of `y` with respect to `x` will be (using zero-based indices, which is what OpenMDAO expects for the `rows` and `cols` arguments): |
| 68 | + |
| 69 | +``` |
| 70 | +y indices: x indices: |
| 71 | +(0, 0) (0, 0) |
| 72 | +(0, 1) (0, 1) |
| 73 | +(0, 2) (0, 2) |
| 74 | +(1, 0) (1, 0) |
| 75 | +(1, 1) (1, 1) |
| 76 | +(1, 2) (1, 2) |
| 77 | +``` |
| 78 | + |
| 79 | +So that table says that the value of `y[0, 0]` depends on `x[0, 0]` only, and the value of `y[1, 0]` depends on `x[1, 0]` only, etc.. |
| 80 | +But OpenMDAO expects flattened indices for the `rows` and `cols` arguments, not multi-dimensional indices. |
| 81 | +So we need to convert the multi-dimensional indices in that table to flattened ones. |
| 82 | +`get_rows_cols` does that for you, but if you wanted to do that by hand, what I usually do is think of an array having the same shape as each input or output, with each entry in the array corresponding to the entry's flat index. |
| 83 | +So for `x` and y, that would be: |
| 84 | + |
| 85 | +``` |
| 86 | +x_flat_indices = |
| 87 | +[0 1 2; |
| 88 | + 3 4 5] |
| 89 | +
|
| 90 | +y_flat_indices = |
| 91 | +[0 1 2; |
| 92 | + 3 4 5] |
| 93 | +``` |
| 94 | + |
| 95 | +(Remember that Python/NumPy arrays use row-major aka C ordering by default.) |
| 96 | +So we can now use those two arrays to translate the `y indices` and `x indices` from multi-dimensional to flat: |
| 97 | + |
| 98 | +``` |
| 99 | +y indices x indices |
| 100 | +multi, flat: multi, flat: |
| 101 | +(0, 0) 0 (0, 0) 0 |
| 102 | +(0, 1) 1 (0, 1) 1 |
| 103 | +(0, 2) 2 (0, 2) 2 |
| 104 | +(1, 0) 3 (1, 0) 3 |
| 105 | +(1, 1) 4 (1, 1) 4 |
| 106 | +(1, 2) 5 (1, 2) 5 |
| 107 | +``` |
| 108 | + |
| 109 | +So the `rows` and `cols` arguments will be |
| 110 | + |
| 111 | +``` |
| 112 | +rows = [0, 1, 2, 3, 4, 5] |
| 113 | +cols = [0, 1, 2, 3, 4, 5] |
| 114 | +``` |
| 115 | + |
| 116 | +where `rows` is the flat non-zero indices for `y`, and `cols` is the flat non-zero indices for `x`. |
| 117 | + |
| 118 | +Now, how do we do this with `get_rows_cols`? |
| 119 | +First we have to assign labels to each dimension of `y` and `x`. |
| 120 | +The labels must be `Symbols`, and can be anything (but I usually use index-y things like `:i`, `:j`, `:k`, etc.). |
| 121 | +We express the sparsity pattern through the choice of labels. |
| 122 | +If we use a label for an output dimension that is also used for an input dimension, then we are saying that, for a given index `i` in the "shared" dimension, the value of the output at that index `i` depends on the value of the input index `i` along the labeled dimension, and no others. |
| 123 | +For example, if we had a one-dimensional `y` that was calculated from a one-dimensional `x` in this way: |
| 124 | + |
| 125 | +``` |
| 126 | +for i in 1:10 |
| 127 | + y[i] = sin(x[i]) |
| 128 | +end |
| 129 | +``` |
| 130 | + |
| 131 | +then we would use the same label for the (single) output and input dimension. |
| 132 | + |
| 133 | +For the present example, we could assign `i` and `j` (say) to the first and second dimensions, respectively, of both `y` and `x`, since `y[i,j]` only depends on `x[i,j]` for all valid `i` and `j`. |
| 134 | +We call these `of_ss` (short for "of subscripts for the output) and `wrt_ss` ("with respect to subscripts"). |
| 135 | + |
| 136 | +```@example shape_by_conn1 |
| 137 | +of_ss = [:i, :j] |
| 138 | +wrt_ss = [:i, :j] |
| 139 | +``` |
| 140 | + |
| 141 | +After deciding on the dimension labels, the only other thing we need to do is create a `Dict` that maps the dimension labels to their sizes: |
| 142 | + |
| 143 | +```@example shape_by_conn1 |
| 144 | +ss_sizes = Dict(:i=>2, :j=>3) |
| 145 | +``` |
| 146 | + |
| 147 | +since, in our example, the first dimension of `x` and `y` has size `2`, and the second, `3`. |
| 148 | + |
| 149 | +Then we pass those three things to `get_rows_cols`, which then returns the `rows` and `cols` we want. |
| 150 | + |
| 151 | +```@example shape_by_conn1 |
| 152 | +rows, cols = OpenMDAOCore.get_rows_cols(; ss_sizes, of_ss, wrt_ss) |
| 153 | +``` |
| 154 | + |
| 155 | +## Back to the Simple Example |
| 156 | +Now, back to the simple example. |
| 157 | +Remember, we're trying to compute `y = 2*x^2 + 1` elementwise for a 2D `x` and `y`. |
| 158 | +The `compute!` method is pretty straight-forward: |
| 159 | + |
| 160 | +```@example shape_by_conn1 |
| 161 | +function OpenMDAOCore.compute!(self::ECompShapeByConn, inputs, outputs) |
| 162 | + x = inputs["x"] |
| 163 | + y = outputs["y"] |
| 164 | + y .= 2 .* x.^2 .+ 1 |
| 165 | + return nothing |
| 166 | +end |
| 167 | +``` |
| 168 | + |
| 169 | +Now, for the `compute_partials!` method, we have to be a bit tricky about the shape of the Jacobian of `y` with respect to `x`. |
| 170 | +The `get_rows_cols` function orders the `rows` and `cols` in such a way that the Jacobian gets allocated by OpenMDAO with shape (`i`, `j`), and is then flattened. |
| 171 | +Since NumPy arrays are row-major ordered, then, we need to reshape the Jacobian in the opposite order, then switch the dimensions. |
| 172 | +This is optional, but makes things easier: |
| 173 | + |
| 174 | +```@example shape_by_conn1 |
| 175 | +function OpenMDAOCore.compute_partials!(self::ECompShapeByConn, inputs, partials) |
| 176 | + x = inputs["x"] |
| 177 | + m, n = size(x) |
| 178 | + # So, with the way I've declared the partials above, OpenMDAO will have |
| 179 | + # created a Numpy array of shape (m, n) and then flattened it. So, to get |
| 180 | + # that to work, I'll need to do this: |
| 181 | + dydx = PermutedDimsArray(reshape(partials["y", "x"], n, m), (2, 1)) |
| 182 | + dydx .= 4 .* x |
| 183 | + return nothing |
| 184 | +end |
| 185 | +``` |
| 186 | + |
| 187 | +## Checking |
| 188 | +Now, let's actually create a `Problem` with the new `Component`, along with an `IndepVarComp` that will actually decide on the size: |
| 189 | + |
| 190 | +```@example shape_by_conn1 |
| 191 | +using OpenMDAO, PythonCall |
| 192 | +
|
| 193 | +m, n = 3, 4 |
| 194 | +p = om.Problem() |
| 195 | +comp = om.IndepVarComp() |
| 196 | +comp.add_output("x", shape=(m, n)) |
| 197 | +p.model.add_subsystem("inputs_comp", comp, promotes_outputs=["x"]) |
| 198 | +
|
| 199 | +ecomp = ECompShapeByConn() |
| 200 | +comp = make_component(ecomp) |
| 201 | +p.model.add_subsystem("ecomp", comp, promotes_inputs=["x"], promotes_outputs=["y"]) |
| 202 | +p.setup(force_alloc_complex=true) |
| 203 | +``` |
| 204 | + |
| 205 | +Now we should be able to check that the output we get is correct: |
| 206 | + |
| 207 | +```@example shape_by_conn1 |
| 208 | +p.set_val("x", 1:m*n) |
| 209 | +p.run_model() |
| 210 | +
|
| 211 | +# Test that the output is what we expect. |
| 212 | +expected = 2 .* PyArray(p.get_val("x")).^2 .+ 1 |
| 213 | +actual = PyArray(p.get_val("y")) |
| 214 | +println("expected = $(expected)") |
| 215 | +println("actual = $(actual)") |
| 216 | +``` |
| 217 | + |
| 218 | +And we can check the derivatives: |
| 219 | + |
| 220 | +```@example shape_by_conn1 |
| 221 | +p.check_partials(method="cs") |
| 222 | +nothing |
| 223 | +``` |
| 224 | + |
| 225 | +Looks good! |
0 commit comments