I have the following code in Pyomo:

from pyomo.environ import *

model = ConcreteModel()
  
model.A = Set(initialize=[1,2,3])
model.B = Set(initialize=['J', 'K'])

model.x = Var(model.A, model.B, bounds=(0, 2))
model.z = Var(model.A, model.B, domain = Binary)

LB = 1
UB = 2

DOMAIN_PTS = [0,0,LB,UB]
RANGE_PTS  = [1,1,0,0]
model.z_constraint = Piecewise(
      model.A, model.B,
      model.z, model.x,
      pw_pts=DOMAIN_PTS,
      pw_repn='INC',
      pw_constr_type = 'EQ',
      f_rule = RANGE_PTS,
      unbounded_domain_var = True)

def objective_rule(model):
    return sum(model.x[a,b] for a in model.A for b in model.B)
model.objective = Objective(rule = objective_rule, sense=minimize)

which creates the following piecewise function:

               / 1    , 0 <= model.x(a,b) <= 0
model.z(a,b) =|   
                0   ,  LB <= model.x(a,b) <= UB

How can I make it so that for each a in model.A, there are unique values for LB and UB, so that the piecewise function becomes:

               / 1    , 0 <= model.x(a,b) <= 0
model.z(a,b) =|   
                0   ,  LB[a] <= model.x(a,b) <= UB[a]

Any help at all would be appreciated.