Sellar

The Sellar problem is a classic problem with one coupled discipline:

\[ \text{minimize }_{x_1, z_1, z_2} f\]

subject to:

\[\begin{split} \begin{aligned} y_1 &\geq 3.6 y_2 &\leq 24 \\ -10 \leq z_1 &\leq 10 \\ 0 \leq z_2 &\leq 10 \\ 0 \leq x_1 &\leq 10 \end{aligned} \end{split}\]

and with

\[\begin{split} \begin{aligned} y_1 &= z_1^2 + z_2 + x_1 - 0.2 y_2\\ y_2 &= \sqrt{y_1} + z_1 + z_2\\ f &= x_1^2 + z_2 + y_1 + \exp(-y_2) \end{aligned} \end{split}\]

We’ll look a bit on how Cottax would handle this internally, but first let’s just look at what you’d do in practice:

import jax
jax.config.update('jax_enable_x64', True)
import jax.numpy as jnp
import cottax.assemblies.jax as cx
from cottax.assemblies.jax import Start
flat = cx.flat

# Declaring
@flat.function("y_1")
def D1(x_1, z_1, z_2, y_2):
    return z_1**2 + z_2 + x_1 - 0.2 * y_2

@flat.function("y_2")
def D2(y_1, z_1, z_2):
    return jnp.sqrt(y_1) + z_1 + z_2

@flat.function("f")
def obj(x_1, z_2, y_1, y_2):
    return x_1**2 + z_2 + y_1 + jnp.exp(-y_2)

@flat.inequality_constraint("g_1")
def C1(y_1):
    return 3.6 - y_1

@flat.inequality_constraint("g_2")
def C2(y_2):
    return y_2 -  24

opt      = flat.optimise(design = {"z_1" : (-10 ,10), "z_2": (0, 10),  "x_1": (0, 10)}, minimise = obj , subject_to = (C1, C2))
declared = flat.declare(D1, D2, obj, C1, C2, opt)

# Running

drivers = (cx.SLSQP(), cx.Newton(), cx.Picard())
mdf      = declared + flat.mdf(opt) + flat.drive(*drivers)

start_val = {Start: dict(z_1 = 5., z_2 = 2., x_1 = 1., y_2 = 0.4)}

out = flat.run(mdf, read_from = start_val)

for name in ("z_1", "z_2", "x_1", "f", "g_1", "g_2"):
    print(f'{name:<3}', out[name])

z_1 2.097366596103213
z_2 3.935152493842969e-14
x_1 0.0
f   3.6184123583244787
g_1 -8.322231792590173e-12
g_2 -20.005266807793525

cx.render_html(mdf, blocks=True, formatter = cx.LatexFormatter())

Breaking it down

The cottax graph itself doesn’t use any implicit conversions and is totally unambiguous. This enables the mathematics to work without having to continuously convert meanings or even have two meanings silently be the same.

Our goal is then to provide the surface interface to the core, which does have these implicit conversions. These surfaces can be swapped if another conversion is desired without affecting any of the core functionality.

This similarly holds for the output side: a visualization in most cases requires a formatter to make the cottax native types legible or transform their representation to nicely-looking latex.

The full representation can be shown below:

Show code cell source

Hide code cell source

import dataclasses

def place(p):    
    keys = ', '.join(repr(k)
                     for k in p.segments)
    return f'{type(p).__name__}({keys})'

def full(x, depth=1):
    pad, inner = '    ' * (depth - 1), '    ' * depth
    if hasattr(x, 'segments'):                        # a VarPath or NodePath, on one line
        return place(x)
    if isinstance(x, tuple):
        if not x: return '()'
        return '(\n' + ''.join(f'{inner}{full(v, depth + 1)},\n' for v in x) + pad + ')'
    if dataclasses.is_dataclass(x) and not isinstance(x, type) and dataclasses.fields(x):
        body = ''.join(f'{inner}{f.name} = {full(getattr(x, f.name), depth + 1)},\n'
                       for f in dataclasses.fields(x))
        return f'{type(x).__name__}(\n{body}{pad})'
    return getattr(x, '__qualname__', None) or repr(x)  # Eq, Min, a function

def show_full(g):
    for n in g.nodes:
        print(f'{place(n)}: {full(g[n])}')
    for inner, outer in g.within.items():
        print(f'{place(inner)} within {place(outer)}')

show_full(mdf.graph)

Show code cell output

Hide code cell output

NodePath(DictKey(key='D1')): ImplementedFunction(
    node = Function(
        reads = (
            VarPath(DictKey(key='x_1')),
            VarPath(DictKey(key='z_1')),
            VarPath(DictKey(key='z_2')),
            VarPath(^copy, DictKey(key='y_2')),
        ),
        owns = (
            VarPath(DictKey(key='y_1')),
        ),
    ),
    implementation = D1.fn,
)
NodePath(DictKey(key='D2')): ImplementedFunction(
    node = Function(
        reads = (
            VarPath(DictKey(key='y_1')),
            VarPath(DictKey(key='z_1')),
            VarPath(DictKey(key='z_2')),
        ),
        owns = (
            VarPath(DictKey(key='y_2')),
        ),
    ),
    implementation = D2.fn,
)
NodePath(DictKey(key='obj')): ImplementedFunction(
    node = Function(
        reads = (
            VarPath(DictKey(key='x_1')),
            VarPath(DictKey(key='z_2')),
            VarPath(DictKey(key='y_1')),
            VarPath(DictKey(key='y_2')),
        ),
        owns = (
            VarPath(DictKey(key='f')),
        ),
    ),
    implementation = obj.fn,
)
NodePath(DictKey(key='C1')): ImplementedFunction(
    node = Function(
        reads = (
            VarPath(DictKey(key='y_1')),
        ),
        owns = (
            VarPath(DictKey(key='g_1')),
        ),
    ),
    implementation = C1.fn,
)
NodePath(DictKey(key='C2')): ImplementedFunction(
    node = Function(
        reads = (
            VarPath(DictKey(key='y_2')),
        ),
        owns = (
            VarPath(DictKey(key='g_2')),
        ),
    ),
    implementation = C2.fn,
)
NodePath(^solve, DictKey(key='f')): ImplementedDriven(
    node = Driven(
        condition = RelationalCondition(
            relations = (
                Relation(
                    lhs = VarPath(^lower, DictKey(key='z_1')),
                    rhs = VarPath(DictKey(key='z_1')),
                    op = Le,
                ),
                Relation(
                    lhs = VarPath(DictKey(key='z_1')),
                    rhs = VarPath(^upper, DictKey(key='z_1')),
                    op = Le,
                ),
                Relation(
                    lhs = None,
                    rhs = VarPath(DictKey(key='z_2')),
                    op = Le,
                ),
                Relation(
                    lhs = VarPath(DictKey(key='z_2')),
                    rhs = VarPath(^upper, DictKey(key='z_2')),
                    op = Le,
                ),
                Relation(
                    lhs = None,
                    rhs = VarPath(DictKey(key='x_1')),
                    op = Le,
                ),
                Relation(
                    lhs = VarPath(DictKey(key='x_1')),
                    rhs = VarPath(^upper, DictKey(key='x_1')),
                    op = Le,
                ),
                Relation(
                    lhs = VarPath(DictKey(key='g_1')),
                    rhs = None,
                    op = Le,
                ),
                Relation(
                    lhs = VarPath(DictKey(key='g_2')),
                    rhs = None,
                    op = Le,
                ),
            ),
            objectives = (
                Objective(
                    place = VarPath(DictKey(key='f')),
                    sense = Min,
                ),
            ),
            unknowns = (
                VarPath(DictKey(key='z_1')),
                VarPath(DictKey(key='z_2')),
                VarPath(DictKey(key='x_1')),
            ),
            parameters = (
                VarPath(^lower, DictKey(key='z_1')),
                VarPath(^upper, DictKey(key='z_1')),
                VarPath(^upper, DictKey(key='z_2')),
                VarPath(^upper, DictKey(key='x_1')),
            ),
        ),
        contract = ConstrainedMinimisation(
            steps = False,
        ),
        data = (
            (
                VarPath(^start, DictKey(key='z_1')),
                VarPath(^start, DictKey(key='z_2')),
                VarPath(^start, DictKey(key='x_1')),
            ),
        ),
        reports = (),
    ),
    implementation = SLSQP(
        tol = 1e-08,
        max_steps = 200,
        strict = True,
        jacobian = jacfwd,
    ),
)
NodePath(^closure, DictKey(key='D1')): ImplementedDriven(
    node = Driven(
        condition = RelationalCondition(
            relations = (
                Relation(
                    lhs = VarPath(^copy, DictKey(key='y_2')),
                    rhs = VarPath(DictKey(key='y_2')),
                    op = Eq,
                ),
            ),
            objectives = (),
            unknowns = (
                VarPath(^copy, DictKey(key='y_2')),
            ),
            parameters = (),
        ),
        contract = RootFinding(
            steps = False,
        ),
        data = (
            (
                VarPath(^start, ^copy, DictKey(key='y_2')),
            ),
        ),
        reports = (),
    ),
    implementation = Newton(
        rtol = 0.0001,
        atol = 0.0001,
        max_steps = 256,
    ),
)
NodePath(^closure, DictKey(key='D1')) within NodePath(^solve, DictKey(key='f'))

This is clearly not easy to write by hand! The surfaces and assemblies that can be imported make a much more convenient way to write this. In fact, the spelling itself is usually shortened a lot to keep it more readable.

A quick explanation of the above:

• VarPath is the variable PyTree, Node-Path. These are different: a node label is not a variable place unless you explicitly convert!

• DictKey is a dictionary key in the datastructure. ^ indicates a name that is created (NOT in the global datastructure by definition!)

• relations are comparisons between left-hand-side and right-hand-sides: None means zero.

• inputs are reads of a function, outputs the outputs of a function. unknowns are the variables a condition variable owns and condition is the condition it needs to satisfy. parameters are variables that affect the answer but are not conditions: data is driver data that is needed for the driver but does not participate in what makes an answer correct (an initial condition!)

• a contract is the driver contract which adds the edges for the initial value (the driver data)

• a implementation refers to the function or some implementation of the driver.

Given values

Also, our philosophy is that we don’t really want any actual values in the graph itself: it must be the structure only. Then, the graph gives a concrete view on what has to be provided to run it: data hidden in the graph obfuscates this. Furthermore, the graph defines the program (through e.g. a JIT cache key), and if values are changing, the compilation has to be redone. [1]

This does lead to some user-friendlyness problems: if you declare an optimizer, you usually give it bounds, and you don’t want to index into some place to give them. The declaration and surface consumers help here. They contain both the graph and a parallel datastructure with those given values and their names in the graph. Additional tags can be used as well (which also have a parallel datastructure).

The MDF Declared object has the following given:

print(mdf.given)

{^lower['z_1']: -10, ^upper['z_1']: 10, ^upper['z_2']: 10, ^upper['x_1']: 10}

And we can automatically do that with function reads as well:

@flat.function("k", reads = {"b" :3})
def f3(a, b ,c):
    ...

f3.given

{^given['k']['b']: 3}

These are given a special place: this wouldn’t be necessary in the flat space, but in hierarchical spaces there may not be a value in the datastructure to express it easily and it is convenient to give.

Then, the above problem requires a set of start values: they are not the same as the variables output (a graph cannot overwrite boundary inputs by definition) and have therefore to be given as well. Again, this is not so convenient in general, because they might have these minted ^ names.

mdf.graph.graph.boundary_inputs

(VarPath(^lower['z_1']),
 VarPath(^upper['z_1']),
 VarPath(^upper['z_2']),
 VarPath(^upper['x_1']),
 VarPath(^start['z_1']),
 VarPath(^start['z_2']),
 VarPath(^start['x_1']),
 VarPath(^start^copy['y_2']))

The assembly run can extract these ^start values and their copies (^start^copy[‘y_2’]) from a datastructure:

flat.valued(mdf, read_from={Start :  dict(z_1 = 5., z_2 = 2., x_1 = 1., y_2 = 0.4)}).values

{^lower['z_1']: -10, ^upper['z_1']: 10, ^upper['z_2']: 10, ^upper['x_1']: 10, ^start['z_1']: 5.0, ^start['z_2']: 2.0, ^start['x_1']: 1.0, ^start^copy['y_2']: 0.4}

---

[1]

In fact, in cottax, this is not necessarily enforced to give you the flexibility to i.e. give some solver setting or hard-coded tolerances, which are non-array-like and therefore hash with it (the above drivers have default tolerances)