while numSteps < minSteps or (rmsGrad > tol and numSteps < maxNumSteps):
# Compute the field energy, with tf watching the variables
with tf.GradientTape() as tape:
energy = lossFn()
vars = [scalarFieldVar, gaugeFieldVar]
# Compute the gradients using automatic differentiation
grads = tape.gradient(energy, vars)
# Postprocess the gauge field gradients so they point in the tangent space
# to SU(2)
grads[1] = FieldTools.projectSu2Gradients(grads[1], gaugeFieldVar)
# Compute max gradient for stopping criterion
gradSq = FieldTools.innerProduct(grads[0], grads[0], tr=True)
gradSq += FieldTools.innerProduct(grads[1], grads[1], tr=True, adj=True)
rmsGrad = tf.sqrt(gradSq)
if (numSteps % printIncrement == 0):
print("Energy after " + str(numSteps) + " iterations: " +\
str(energy.numpy()))
print("RSS gradient after " + str(numSteps) + " iterations: " +\
str(rmsGrad.numpy()))
# Perform the gradient descent step
opt.apply_gradients(zip(grads, vars))
numSteps += 1
# Postprocess the fields to avoid drift away from SU(2)/its Lie algebra
opt = tf.keras.optimizers.SGD(learning_rate=0.01*args.gaugeCoupling*args.vev)
while numSteps < minSteps or (rssGrad < rssGradOld and numSteps < maxSteps and rssGrad > tol):
# Compute the field energy, with tf watching the variables
with tf.GradientTape() as tape:
energy = lossFn()
vars = [scalarFieldVar, gaugeFieldVar]
# Compute the gradients using automatic differentiation
grads = tape.gradient(energy, vars)
# Postprocess the gauge field gradients
grads = theory.processGradients(grads, vars)
# Compute RSS gradient for stopping criterion
gradSq = FieldTools.innerProduct(grads[0], grads[0], tr=True)
gradSq += FieldTools.innerProduct(grads[1], grads[1], tr=True, adj=True)
rssGradOld = rssGrad
rssGrad = tf.math.sqrt(gradSq)
# rssGrad = tf.reduce_max(tf.abs(grads[1]))
if (numSteps % printIncrement == 0):
print("Energy after " + str(numSteps) + " iterations: " +\
str(energy.numpy()))
print("RSS gradient after " + str(numSteps) + " iterations: " +\
str(rssGrad.numpy()))
# Perform the gradient descent step
opt.apply_gradients(zip(grads, vars))
numSteps += 1