Adding custom variables to PyPSA objective function without requiring a change to optimize.py in base pypsa code?

I am modelling a network in pypsa with constraints that have statistically calculated RHS, meaning that there is always a chance of an infeasible solution (e.g. X > 10 & X < 9)

To get around this I add a penalty term into the objective function, example given below. However I have had to modify optimize.py as follows:

def assign_solution(n: Network) -> None:
"""
Map solution to network components.
"""
m = n.model
sns = n.model.parameters.snapshots.to_index()

for name, variable in m.variables.items():
    sol = variable.solution
    if name == "objective_constant":
        continue
    **if name[:7].lower() == "penalty":
        continue**
    try:
        c, attr = name.split("-", 1)
        df = sol.to_pandas()
    except ValueError:
        continue

Where the added lines are

    if name[:7].lower() == "penalty":
        continue

If I don't make this change pypsa fails with a KeyError because the function assign_solution which 'Maps solution to network components' can't find a component for the new objective variable.

Is there some way to add these penalty variables to some inert component, or can we get something like these two lines added to the pypsa repo, to help anybody else using constraint relaxation. Or is there another solution out there?

Toy example of constraint relaxation, requiring above modifications to optimize.py

import pypsa

def build_model():
    # Create toy network
    n = pypsa.Network()
    n.add("Carrier", "AC")
    n.add("Bus", "Bus 0",carrier = 'AC')
    n.add("Bus", "Bus 1",carrier = 'AC')
    n.add("Line","Line 01", bus0 = "Bus 0", bus1 = "Bus 1", x=0.1, s_nom=50, carrier =     'AC', r = 1)

    n.add("Generator", "Gen 0",
                bus=f"Bus 0",
                p_nom=100,
                carrier = 'AC',
                marginal_cost=10,
        )

    n.add("Load", "Load 0",
                bus="Bus 1",
                p_set=50,
                carrier = 'AC',
        )
    n.optimize.create_model()

return n

n = build_model()
n.optimize.solve_model()

This network of a 100MW generator connected to a 50MW by 1 line solves, no problem

However if I add a constraint limiting that line to 40MW, it does not

# Set constraint restricting line flow to 40MW or less (infeasible, as load is 50MW)
n = build_model()
m = n.model
line_flow = m.variables["Line-s"].sel({'Line':'Line 01'})
m.add_constraints(line_flow <= 40)
n.optimize.solve_model()

Solver message: infeasible ('warning', 'infeasible')

So I can add constraint relaxation

# Add constraint relaxation
n = build_model()
m = n.model
penalty = m.add_variables(lower = 0, name = f"Penalty_varible", coords = [n.snapshots])
extra_objective_term = (penalty * 1000).sum() # Add weighting of 1000 to the penalty term
n.model.objective = m.objective + extra_objective_term 
m.add_constraints(line_flow <= penalty + 40)
n.optimize.solve_model()

n.model.solution

Which solves and gives penalty = 10, as you'd expect

Is there any way to make this compatible with PyPSA without having to change the library code?

Cheers