Skip to content

Add constraints scaling to ~ order unity #4718

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
timothy-nunn wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add constraints scaling to ~ order unity #4718

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
43a25df
Add in option to scale constraints
timothy-nunn Oct 28, 2026
68cfdc4
Add more scales
timothy-nunn Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 68 additions & 26 deletions process/core/solver/constraints.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,12 @@ class ConstraintRegistration:
"""The type of constraint (<=, >=, ==). Only used for writing output diagnostics,
this does not impact the calculations.
"""
scale: float
"""The scale of the constraint.

I.e. what the constraint is multiplied by to convert it to be order unity when passed
to the solver.
"""


@dataclass
Expand Down Expand Up @@ -100,7 +106,12 @@ def constraint_ids(cls):

@classmethod
def register_constraint(
cls, name: Hashable, units: str, symbol: ConstraintSymbolType
cls,
name: Hashable,
units: str,
symbol: ConstraintSymbolType,
*,
scale: float = 1.0,
) -> Callable[[], Callable[[], ConstraintResult]]:
"""A decorator to add a constraint equation with metadata to the registry.

Expand All @@ -115,14 +126,17 @@ def register_constraint(
the units of the constraint written to the output files
symbol : str
the symbol of the constraint written to the output files
scale : float
a factor that the constraint value and bound is multiplied by before being
sent to the solver

"""

def wrapper(func: Callable[[], ConstraintResult]):
if name in cls._constraint_registry:
raise ValueError(f"Constraint {name} already exists.")
cls._constraint_registry[name] = ConstraintRegistration(
name, func, units, symbol
name, func, units, symbol, scale
)

return func
Expand All @@ -147,7 +161,9 @@ def get_constraint(cls, name: Hashable):
return cls._constraint_registry.get(name)

@classmethod
def evaluate_constraint(cls, name: Hashable, data: DataStructure):
def evaluate_constraint(
cls, name: Hashable, data: DataStructure, *, with_scale: bool = True
):
"""Evaluates a constraint with a given name.

Parameters
Expand All @@ -174,9 +190,27 @@ def evaluate_constraint(cls, name: Hashable, data: DataStructure):
error_msg = f"Constraint '{name}' cannot be found."
raise ProcessError(error_msg)

if not with_scale:
registration = ConstraintRegistration(**{
**asdict(registration),
"scale": 1.0,
})

return registration.constraint_equation(registration, data)


def _apply_scale(func):
"""Scales the value and bound before passing them on to the constraint function."""

def _wrapper(value: float, bound: float, registration: ConstraintRegistration):
v = value * registration.scale
b = bound * registration.scale
return func(v, b, registration)

return _wrapper


@_apply_scale
def leq(value: float, bound: float, registration: ConstraintRegistration):
"""The equation `value <= bound`."""
residual = value - bound
Expand All @@ -190,6 +224,7 @@ def leq(value: float, bound: float, registration: ConstraintRegistration):
)


@_apply_scale
def geq(value: float, bound: float, registration: ConstraintRegistration):
"""The equation `value >= bound`."""
residual = bound - value
Expand All @@ -203,6 +238,7 @@ def geq(value: float, bound: float, registration: ConstraintRegistration):
)


@_apply_scale
def eq(value: float, bound: float, registration: ConstraintRegistration):
"""The equation `value = bound`."""
residual = value - bound
Expand Down Expand Up @@ -253,7 +289,7 @@ def constraint_equation_1(constraint_registration, data):
)


@ConstraintManager.register_constraint(2, "MW", "=")
@ConstraintManager.register_constraint(2, "MW", "=", scale=1e-2)
def constraint_equation_2(constraint_registration, data):
"""
Global plasma power balance equation [MW]
Expand Down Expand Up @@ -407,7 +443,7 @@ def constraint_equation_4(constraint_registration, data):
return eq(pnumerator, pdenom, constraint_registration)


@ConstraintManager.register_constraint(5, "/m³", "<=")
@ConstraintManager.register_constraint(5, "/m³", "<=", scale=1e-20)
def constraint_equation_5(constraint_registration, data):
"""Equation for electron density upper limit

Expand Down Expand Up @@ -557,7 +593,7 @@ def constraint_equation_12(constraint_registration, data):
)


@ConstraintManager.register_constraint(13, "sec", ">=")
@ConstraintManager.register_constraint(13, "sec", ">=", scale=1e-3)
def constraint_equation_13(constraint_registration, data):
"""Equation for burn time lower limit

Expand Down Expand Up @@ -608,7 +644,7 @@ def constraint_equation_15(constraint_registration, data):
)


@ConstraintManager.register_constraint(16, "MW", ">=")
@ConstraintManager.register_constraint(16, "MW", ">=", scale=1e-3)
def constraint_equation_16(constraint_registration, data):
"""Equation for net electric power lower limit

Expand Down Expand Up @@ -752,7 +788,7 @@ def constraint_equation_23(constraint_registration, data):
)


@ConstraintManager.register_constraint(24, "", "<=")
@ConstraintManager.register_constraint(24, "", "<=", scale=1e2)
def constraint_equation_24(constraint_registration, data):
"""Equation for beta upper limit

Expand Down Expand Up @@ -914,7 +950,7 @@ def constraint_equation_30(constraint_registration, data):
)


@ConstraintManager.register_constraint(31, "Pa", "<=")
@ConstraintManager.register_constraint(31, "Pa", "<=", scale=1e-6)
def constraint_equation_31(constraint_registration, data):
"""Equation for TF coil case stress upper limit (SCTF)

Expand All @@ -930,7 +966,7 @@ def constraint_equation_31(constraint_registration, data):
)


@ConstraintManager.register_constraint(32, "Pa", "<=")
@ConstraintManager.register_constraint(32, "Pa", "<=", scale=1e-6)
def constraint_equation_32(constraint_registration, data):
"""Equation for TF coil conduit stress upper limit (SCTF)

Expand All @@ -946,7 +982,7 @@ def constraint_equation_32(constraint_registration, data):
)


@ConstraintManager.register_constraint(33, "A/m²", "<=")
@ConstraintManager.register_constraint(33, "A/m²", "<=", scale=1e-6)
def constraint_equation_33(constraint_registration, data):
"""Equation for TF coil operating/critical J upper limit (SCTF)

Expand Down Expand Up @@ -980,7 +1016,7 @@ def constraint_equation_34(constraint_registration, data):
)


@ConstraintManager.register_constraint(35, "A/m²", "<=")
@ConstraintManager.register_constraint(35, "A/m²", "<=", scale=1e-6)
def constraint_equation_35(constraint_registration, data):
"""Equation for TF coil J_wp/J_prot upper limit (SCTF)

Expand Down Expand Up @@ -1425,7 +1461,7 @@ def constraint_equation_64(constraint_registration, data):
)


@ConstraintManager.register_constraint(65, "Pa", "<=")
@ConstraintManager.register_constraint(65, "Pa", "<=", scale=1e-6)
def constraint_equation_65(constraint_registration, data):
"""Upper limit on stress of the vacuum vessel that occurs when the TF coil quenches.

Expand Down Expand Up @@ -1513,7 +1549,7 @@ def constraint_equation_68(constraint_registration, data):
)


@ConstraintManager.register_constraint(72, "Pa", "<=")
@ConstraintManager.register_constraint(72, "Pa", "<=", scale=1e-6)
def constraint_equation_72(constraint_registration, data):
"""Upper limit on Central Solenoid Tresca yield stress

Expand Down Expand Up @@ -1721,7 +1757,7 @@ def constraint_equation_80(constraint_registration, data):
)


@ConstraintManager.register_constraint(81, "/m³", ">=")
@ConstraintManager.register_constraint(81, "/m³", ">=", scale=1e-20)
def constraint_equation_81(constraint_registration, data):
"""Lower limit to ensure central density is larger that the pedestal one

Expand Down Expand Up @@ -1876,7 +1912,7 @@ def constraint_equation_89(constraint_registration, data):
)


@ConstraintManager.register_constraint(90, "", ">=")
@ConstraintManager.register_constraint(90, "", ">=", scale=1e-3)
def constraint_equation_90(constraint_registration, data):
"""Lower limit for CS coil stress load cycles

Expand Down Expand Up @@ -1939,24 +1975,21 @@ def constraint_equation_92(constraint_registration, data):
)


@ConstraintManager.register_constraint(93, "Z.particles/s", "=")
@ConstraintManager.register_constraint(93, "particles/s", "=", scale=1e-21)
def constraint_equation_93(constraint_registration, data):
"""Fuel-ion equilibrium constraint.

Requires the plasma fuelling rate to equal the fuelling rate required
for fuel-ion equilibrium.

Note the units of the constraint residual (not either variable)
is zetta-particles per second.
"""
return eq(
data.physics.molflow_plasma_fuelling / 10**21,
data.physics.molflow_plasma_fuelling_equilibrium / 10**21,
data.physics.molflow_plasma_fuelling,
data.physics.molflow_plasma_fuelling_equilibrium,
constraint_registration,
)


def constraint_eqns(m: int, ieqn: int, data: DataStructure):
def constraint_eqns(m: int, ieqn: int, data: DataStructure, *, with_scale: bool = True):
"""Evaluates the constraints given the current state of PROCESS.

Parameters
Expand All @@ -1969,6 +2002,8 @@ def constraint_eqns(m: int, ieqn: int, data: DataStructure):
data:
DataStructure object to provide data to evaluate the constraints
NOTE: this is only for getting data, not setting it
with_scale:
If False, the data is not multiplied by the scale before the calculation.

Raises
------
Expand All @@ -1986,7 +2021,9 @@ def constraint_eqns(m: int, ieqn: int, data: DataStructure):

for i in range(i1, i2):
constraint_id = data.numerics.icc[i]
result = ConstraintManager.evaluate_constraint(constraint_id, data)
result = ConstraintManager.evaluate_constraint(
constraint_id, data, with_scale=with_scale
)

tmp_cc, tmp_con, tmp_err = (
result.normalised_residual,
Expand Down Expand Up @@ -2024,7 +2061,10 @@ def constraints_output(data: DataStructure, solver_name: str):
)

con1, con2, err, _, lab = constraint_eqns(
nums.n_equality_constraints + nums.n_inequality_constraints, -1, data
nums.n_equality_constraints + nums.n_inequality_constraints,
-1,
data,
with_scale=False,
)

# Write equality constraints to mfile
Expand Down Expand Up @@ -2085,7 +2125,9 @@ def constraints_output(data: DataStructure, solver_name: str):
nums.n_equality_constraints + nums.n_inequality_constraints,
):
name = nums.lablcc[nums.icc[i] - 1]
constraint = ConstraintManager.evaluate_constraint(int(nums.icc[i]), data)
constraint = ConstraintManager.evaluate_constraint(
int(nums.icc[i]), data, with_scale=False
)

inequality_constraint_table.append([
name,
Expand Down
Loading