Function Metadata API

Function Metadata API#

Using ExplicitFuncComp, you can turn a python function into a fully functioning OpenMDAO component. However, in order to do that it’s sometimes necessary to attach additional metadata to the function so that OpenMDAO can be informed of things like variable units and shapes, and partial derivative information. Metadata can be attached to a function using the function metadata API. It works by wrapping the function in a callable object that can store the metadata appropriately.

Function wrapping#

We wrap a function using the omf.wrap function, for example:

import openmdao.func_api as omf
import numpy as np

def func(a):
    x = a * 2.

f = omf.wrap(func) 

omf.wrap returns an instance of the OMWrappedFunc class that can store various metadata needed by OpenMDAO. All of the metadata setting functions called on that instance return the instance itself so they can be stacked together. For example:

f = omf.wrap(func).add_input('a', shape=5).add_output('x', shape=5, units='m')

Also, if you need to make many calls to set metadata on the wrapped function, you can stack the calls vertically, but this will only work if you wrap the entire righ-hand-side expression in parentheses so that python will treat it all as a single expression. For example:

f = (omf.wrap(func)
        .defaults(units='m')
        .add_input('a', shape=5)
        .add_output('x', shape=5))

If stacking isn’t desired, the methods can just be called in the usual way, for example:

f = omf.wrap(func)
f.defaults(units='m')
f.add_input('a', shape=5)
f.add_output('x', shape=5)

<openmdao.func_api.OMWrappedFunc at 0x7fd6e88e38d0>

Variable metadata#

Setting the metadata for a single variable#

OpenMDAO needs to know a variable’s shape, initial value, and optionally other things like units.
This information can be specified using the add_input and add_output methods. For example:

def func(x):
    y = x.dot(np.random.random(2))
    return y

f = (omf.wrap(func)
        .add_input('x', shape=(2,2))
        .add_output('y', shape=2))

Setting metadata for option variables#

A function may have additional non-float or non-float ndarray arguments that, at least in the OpenMDAO context, will be treated as component options that don’t change during a given model execution. These can be specified using the declare_option method. For example:

def func(x, opt):
    if opt == 1:
        y = x.dot(np.random.random(2))
    elif opt == 2:
        y = x[:, 1] * 2.
    elif opt == 3:
        y = x[1, :] * 3.
    return y

f = (omf.wrap(func)
        .add_input('x', shape=(2,2))
        .declare_option('opt', values=[1, 2, 3])
        .add_output('y', shape=2))

The arguments that are passable to declare_option are the same as those that are allowed when declaring option variables in an OpenMDAO component using the OptionsDictionary declare method.

Setting metadata for multiple variables#

Using the add_inputs and add_outputs methods you can specify metadata for multiple variables in the same call. For example:

def func(a, b):
    return a.dot(b), a[:,0] * b * b

f = (omf.wrap(func)
        .add_inputs(a={'shape': (2,2), 'units': 'm'}, b={'shape': 2, 'units': 'm'})
        .add_outputs(x={'shape': 2, 'units': 'm**2'}, y={'shape': 2, 'units': 'm**3'}))

Getting the metadata#

Variable metadata is retrieved from the wrapped function by calling the get_input_meta and get_output_meta methods. Each function returns an iterator over (name, metadata_dict) tuples, one for each input or output variable respectively. For example, the following code snippet will print the name and shape of each output variable.

for name, meta in f.get_output_meta():
    print(name, meta['shape'])

x (2,)
y (2,)

Setting function default metadata#

Some metadata will be the same for all, or at least most of the variables within a given function, so we want to be able to specify those defaults easily without too much boilerplate. That’s the purpose of the defaults method. For example:

def func(a, b, c):
    d = a * b * c
    return d

f = omf.wrap(func).defaults(shape=4, units='m')

Any metadata that is specific to a particular variable will override any defaults specified in defaults. For example:

def func(a, b, c=np.ones(3)):  # shape of c is 3 which overrides the `defaults` shape of 4
    d = a * b
    e = c * 1.5
    return d, e

f = omf.wrap(func).defaults(shape=4, units='m')

Assumed default values#

In order to stay consistent with OpenMDAO’s default value policy, we assume the same default behavior for functions, so if no shape or default value is supplied for a function variable, we assume that is has the value 1.0. If the shape is provided and either the default value is not provided or is provided as a scalar value, then the assumed default value will be np.ones(shape) * scalar_value, where scalar_value is 1.0 if not specified. If shape is provided along with a non-scalar default value that has a different shape, then an exception will be raised.

Variable names#

Setting variable names#

We don’t need to set input names because the function can always be inspected for those, but we do need to associate output names with function return values. Those return values, if they are simple variables, for example, return x, y, will give us the output variable names we need.
But in those cases where the function returns expressions rather than simple variables, we need another way to specify what the names of those output variables should be. The output_names method provides a concise way to do this, for example:

def func(a, b, c):
    return a * b * c, a * b -c  # two return values that don't have simple names

f = omf.wrap(func).output_names('d', 'e')  # name of return values are 'd' and 'e'

If we have metadata we need to supply for the outputs, we could instead just use the add_outputs method mentioned earlier, for example:

def func(a, b, c):
    return a * b * c, a * b -c  # two return values that don't have simple names

# names of return values are 'd' and 'e'. 
f = omf.wrap(func).add_outputs(d={'units': 'm'}, e={'units': 'ft'})

As mentioned above, if the function’s return values are simple variable names, we don’t need to specify the output names because we can determine them by inspecting the function, e.g.,

def func(a, b, c):
    d = a * b * c
    e = a * b -c
    return d, e  # we know from inspection that the output names are 'd' and 'e'

Note that in the function above, we didn’t have to wrap it at all. This is possible because we can inspect the source code to determine the output names and we assume the default value of all inputs and outputs is 1.0. If any inputs or outputs of a function have any non-default metadata, e.g., val, units, shape, etc., then that function would have to be wrapped and those metadata values would have to be specified. Also, if we plan to compute derivatives for the function then we would need to specify which partials are nonzero using the declare_partials method.

If one or more output names are not specified and cannot be determined by inspection, then they must be specified using add_output calls. The number of add_output calls corresponding to unnamed return values must match the total number of unnamed return values, and they will be matched to those return values in the order that they are called. Any call to add_output with an output name that corresponds to one already specified can occur in any order. In the example below, there are two return values and neither output name is specified, so two calls to add_output are needed.

def func(x):
    return x.dot(np.random.random(2)), x*1.5  # 2 return values and we can't infer the names
f = (omf.wrap(func)
        .add_input('x', shape=(2,2))
        .add_output('y', shape=2)       # 'y' is the name of the first return value
        .add_output('z', shape=(2,2)))  # 'z' is the name of the second return value

In the example above, the output names would be assumed to be ['y', 'z'].

Getting variable names#

Lists of input names and output names can be retrieved by calling get_input_names and get_output_names respectively, e.g.,

print('input names =', list(f.get_input_names()))
print('output names = ', list(f.get_output_names()))

input names = ['x']
output names =  ['y', 'z']

Partial derivatives#

Setting partial derivative information#

Metadata that will help OpenMDAO to compute partial derivatives for the function can be defined using the declare_partials and declare_coloring methods. For example:

def func(x, y, z=3): 
    foo = np.log(z)/(3*x+2*y)
    bar = 2*x+y
    return foo, bar

f = (omf.wrap(func)
        .declare_partials(of='*', wrt='*', method='cs')
        .declare_coloring(wrt='*', method='cs')
        .defaults(shape=4))

The arguments for the declare_partials and declare_coloring methods match those of the same methods on Component. Multiple calls can be made to declare_partials to set up different partials, but declare_coloring may only be called once.

Note that all nonzero partial derivatives must be declared or OpenMDAO will assume they are zero.

Getting partial derivative information#

The arguments passed to the declare_partials and declare_coloring methods can be retrieved using the get_declare_partials and get_declare_coloring methods respectively. Each of these returns a list where each entry is the keyword args dict from each call, in the order that they where called.

print(f.get_declare_partials())  # returns a list of args dicts for multiple calls
print(f.get_declare_coloring())   # returns args dict for a single call to declare_coloring

[{'method': 'cs', 'of': '*', 'wrt': '*'}]
{'method': 'cs', 'wrt': '*'}