What is the proper way to apply Dong's 2021 JCP periodic BC method in DeepXDE?

[ululi1970]

I am reading Dong and Ni's 2021 JCP paper on how to implement exactly periodic BC with NN. The idea is to map the input coordinates (x,y,...) to (sin(kx),cos(kx),sin(qy),cos(qy),...) and then pass this to the NN. Would this snippet work

# assume 2D problem and tensorflow backend
K = [2*tf.pi/px,2*tf.pi/py]
def Periodic_transform(x):
    
     return tf.concat(tf.sin(K[0]*x[:,0:1]),tf.cos(K[0]*x[:,0:1]),tf.sin(K[1]*x[:,1:2]),tf.cos(K[1]*x[:,1:2]), axis= 1)

net = dde.nn.fnn([4] + [width]*depth + [n_output], "tanh", "Glorot normal")
# the input layer of the net is twice the dimension of the input space correct?
net.apply_feature_transform(Periodic_transform)

# etc