Parameter constraints with LinearConstraint#

This notebook is a runnable companion to the Constraints user guide. It shows how to declare a LinearConstraint, what verify_constraints and enforce_constraints actually do, and how the residual parameterization shows up in autograd Jacobians. The final section uses the production GL06INSOUT model to confirm that the constraint system makes every derived M1 parameter invisible to autograd, and discusses how to read the numerical Jacobian comparison on the free parameters.

import macrostat
import torch

print('macrostat version:', macrostat.__version__)
print('torch version:    ', torch.__version__)
macrostat version: 0.5.7.post1.dev7+g35c4d7655.d20260409
torch version:     2.11.0+cu130

A minimal Parameters subclass#

The smallest demonstration is a three-parameter group whose values must sum to one. We pick the third parameter as the residual.

from macrostat.core.parameters import Parameters
from macrostat.core.constraints import LinearConstraint


class ToyParameters(Parameters):
    def get_default_parameters(self):
        return {
            'share_a': {'value': 0.4, 'lower bound': 0.0,
                        'upper bound': 1.0, 'notation': 'a', 'unit': '.'},
            'share_b': {'value': 0.3, 'lower bound': 0.0,
                        'upper bound': 1.0, 'notation': 'b', 'unit': '.'},
            'share_c': {'value': 0.3, 'lower bound': 0.0,
                        'upper bound': 1.0, 'notation': 'c', 'unit': '.'},
        }

    def get_constraints(self):
        return (
            LinearConstraint(
                param_names=('share_a', 'share_b', 'share_c'),
                target=1.0,
            ),
        )

What happens when defaults violate the constraint?#

Before the success case, deliberately give the toy parameters a set of defaults that do not sum to the target. The verifier emits a warning at construction time, and enforce_constraints then adjusts the residual parameter so the identity holds.

import logging
logging.basicConfig(level=logging.WARNING, force=True)

class ViolatingToy(ToyParameters):
    def get_default_parameters(self):
        defaults = super().get_default_parameters()
        # These sum to 0.9, not 1.0
        defaults['share_a']['value'] = 0.5
        defaults['share_b']['value'] = 0.3
        defaults['share_c']['value'] = 0.1
        return defaults

p_bad = ViolatingToy()
print('after init:')
for k in ('share_a', 'share_b', 'share_c'):
    print(f'  {k} = {p_bad.values[k]["value"]}')
print('sum =', sum(p_bad.values[k]['value'] for k in
                  ('share_a', 'share_b', 'share_c')))
WARNING:macrostat.core.parameters:Constraint violation in defaults: sum(('share_a', 'share_b', 'share_c')) = 0.90000000, target = 1.00000000. Derived parameter 'share_c' will be adjusted.
after init:
  share_a = 0.5
  share_b = 0.3
  share_c = 0.19999999999999996
sum = 1.0

The warning fired during construction; the residual share_c was then overwritten from 0.1 to 0.2 so the group sums to 1.0. The free parameters share_a and share_b were left untouched.

The happy path: free parameters move, derived parameter follows#

Now use the consistent defaults from ToyParameters. Mutate one of the free parameters and re-run enforce_constraints to see the residual track.

p = ToyParameters()
print('initial values:')
for k in ('share_a', 'share_b', 'share_c'):
    print(f'  {k} = {p.values[k]["value"]}')

p.values['share_a']['value'] = 0.55
p.enforce_constraints()
print()
print('after share_a := 0.55 and enforce_constraints():')
for k in ('share_a', 'share_b', 'share_c'):
    print(f'  {k} = {p.values[k]["value"]}')
initial values:
  share_a = 0.4
  share_b = 0.3
  share_c = 0.30000000000000004

after share_a := 0.55 and enforce_constraints():
  share_a = 0.55
  share_b = 0.3
  share_c = 0.1499999999999999
print('all parameters:', list(p.values.keys()))
print('free parameters:', p.get_free_param_names())
all parameters: ['share_a', 'share_b', 'share_c']
free parameters: ['share_a', 'share_b']

Autograd consequence: derived parameter has zero gradient#

Because share_c = 1 - share_a - share_b is encoded as a tensor operation, autograd never treats share_c as a free leaf. Any downstream loss has zero gradient with respect to it; the sensitivity is absorbed into share_a and share_b.

p = ToyParameters()

# Build leaf tensors for the FREE parameters; the residual is then
# derived from them so it is not a leaf and has no gradient of its own.
# Initialise the residual slot as a zero placeholder of the same dtype;
# LinearConstraint.apply will overwrite it through the resolver.
free_names = ('share_a', 'share_b')
tensors = {
    name: torch.tensor(p.values[name]['value'], requires_grad=True)
    for name in free_names
}
tensors['share_c'] = torch.tensor(0.0)

# The step-time apply() path wants a ConstraintResolver. Build one
# against the Parameters instance; for this scalar toy model every
# slot resolves to its own tensor key with no index.
resolver = p.get_constraint_resolver()
for c in p.get_constraints():
    c.apply(tensors, resolver)  # writes share_c = 1 - share_a - share_b

loss = (
    2.0 * tensors['share_a']
    + 3.0 * tensors['share_b']
    + 4.0 * tensors['share_c']
)
loss.backward()

for name in ('share_a', 'share_b'):
    print(f'  d(loss)/d({name}) = {tensors[name].grad.item():+.3f}')
share_c = tensors['share_c']
print(f'  share_c is derived: is_leaf={share_c.is_leaf}, '
      f'requires_grad={share_c.requires_grad}')

The free parameters report 2 - 4 = -2 and 3 - 4 = -1 because the residual coefficient 4 flows back through share_c = 1 - share_a - share_b with d(share_c)/d(share_a) = -1. The derived parameter share_c is a function of the leaves and so has no gradient of its own. This is the residual parameterization in action.

Real-world example: GL06INSOUT#

The Godley-Lavoie INSOUT model has a 4-asset Tobin portfolio with M1 cash as the buffer-stock residual. Five LinearConstraint instances cover the constants and the four rate-sensitivity columns.

The constraint system’s guarantee is narrow and mechanical: every derived M1 parameter is a tensor function of the free parameters, not a leaf, so autograd sees exactly zero sensitivity with respect to it. We check that first. Numerical-vs-autograd agreement on the free parameters is a separate question that depends on the rest of the model’s numerical behaviour, which we discuss after the comparison.

from macrostat.models import get_model

model = get_model('GL06INSOUT')()
constraints = model.parameters.get_constraints()
print(f'GL06INSOUT declares {len(constraints)} constraints:')
for i, c in enumerate(constraints, 1):
    print(f'  {i}. target={c.target}')
    print(f'     free   : {list(c.free_params)}')
    print(f'     derived: {c.derived_param}')
GL06INSOUT declares 5 constraints:
  1. target=1.0
     free   : ['WealthShareM2_Constant', 'WealthShareBills_Constant', 'WealthShareBonds_Constant']
     derived: WealthShareM1_Constant
  2. target=0.0
     free   : ['WealthShareM2_DepositRate', 'WealthShareBills_DepositRate', 'WealthShareBonds_DepositRate']
     derived: WealthShareM1_DepositRate
  3. target=0.0
     free   : ['WealthShareM2_BillRate', 'WealthShareBills_BillRate', 'WealthShareBonds_BillRate']
     derived: WealthShareM1_BillRate
  4. target=0.0
     free   : ['WealthShareM2_BondYield', 'WealthShareBills_BondYield', 'WealthShareBonds_BondYield']
     derived: WealthShareM1_BondYield
  5. target=0.0
     free   : ['WealthShareM2_Income', 'WealthShareBills_Income', 'WealthShareBonds_Income']
     derived: WealthShareM1_Income
import sys, os
# Make the doc helper module importable from the notebook
sys.path.insert(0, os.path.abspath('.'))
from _doc_helpers import mean_nominal_output

from macrostat.diff import (
    JacobianAutograd,
    JacobianNumerical,
    compare_jacobian_dicts,
)

# Hand-picked subset: a few free parameters from the Tobin matrix
# (where the constraint system matters) plus two unrelated parameters
# as a sanity check. A full sweep over all 42 free parameters takes
# minutes; for a worked example, five is enough to demonstrate the
# pattern.
sample_params = [
    'WealthShareM2_Constant',
    'WealthShareBills_Constant',
    'WealthShareBonds_Constant',
    'TaxRate',
    'InventorySalesRatioBaseline',
]

jac_auto = JacobianAutograd(model).compute(
    loss_fn=mean_nominal_output, mode='rev',
)
jac_num = JacobianNumerical(model).compute(
    loss_fn=mean_nominal_output, mode='central', num_workers=1,
    param_names=sample_params,
)

jac_auto_subset = {k: jac_auto[k] for k in sample_params}

report = compare_jacobian_dicts(
    jac_auto_subset, jac_num,
    method_a='autograd', method_b='numerical (log, eps=1e-3)',
    rtol=1e-3, atol=1e-6,
)
print(report.summary())
Jacobian comparison: autograd vs numerical (log, eps=1e-3)
Parameters compared: 5
Overall max abs diff: 1.899e+00  max rel diff: 2.077e-01

Parameter                                   max_abs   mean_abs    max_rel   mean_rel        close nan/inf
---------------------------------------------------------------------------------------------------------
WealthShareBills_Constant                 1.899e+00  1.899e+00  2.077e-01  2.077e-01          0/1      no
WealthShareM2_Constant                    5.841e-02  5.841e-02  5.841e-02  5.841e-02          0/1      no
WealthShareBonds_Constant                 1.534e-01  1.534e-01  7.729e-03  7.729e-03          0/1      no
InventorySalesRatioBaseline               9.693e-02  9.693e-02  3.941e-03  3.941e-03          0/1      no
TaxRate                                   3.356e-01  3.356e-01  1.868e-04  1.868e-04          1/1      no

TaxRate agrees to four significant digits. The two unrelated parameters and WealthShareBonds_Constant agree to roughly 1%; the two other Tobin constants show disagreements of 6% and 21%. This residual disagreement is not a constraint-system bug. GL06INSOUT’s portfolio_switches method uses a step-function torch.where to select between the M1- and M2-dominated regimes, and any parameter that moves the economy near that boundary (which the Tobin constants do by shifting the baseline portfolio composition) straddles the discontinuity under finite differences. Autograd reports zero gradient at the switch; numerical differences see a finite jump. See the “Expected disagreements” section of Differentiability and Jacobian Tools and the worked example in Jacobian Diagnostics.

The next cell runs the check that is the constraint-system guarantee: every derived M1 parameter should have exactly zero autograd gradient.

# Confirm that every derived M1 parameter has exactly zero autograd gradient
derived = [c.derived_param for c in constraints]
for name in derived:
    g = jac_auto[name]
    print(f'  {name:<35s} max|grad| = {g.abs().max().item():.3e}')
  WealthShareM1_Constant              max|grad| = 0.000e+00
  WealthShareM1_DepositRate           max|grad| = 0.000e+00
  WealthShareM1_BillRate              max|grad| = 0.000e+00
  WealthShareM1_BondYield             max|grad| = 0.000e+00
  WealthShareM1_Income                max|grad| = 0.000e+00

Closing notes#

  • The constraint system is opt-in: existing models that do not override get_constraints are unaffected.

  • LinearConstraint.apply runs inside Behavior.apply_parameter_shocks at every step, fetched fresh against the current Parameters instance, so scenario shocks to free parameters automatically propagate to the residual. The init-time adjustment happens via Parameters.enforce_constraints.

  • For the API and the verification rules, see Constraints.

  • For diagnosing autograd vs numerical disagreements that are not explained by constraints, see Jacobian Diagnostics.