About the knowledge control model

Hello, I noticed that the knowledge control imposes a restriction on the predicted results, the knowledge control added to the N-body MNIST dataset is the law of conservation of energy, and the knowledge control added to the SEVIR dataset is the expected precipitation intensity, good idea! I want to use the knowledge control on other datasets to make model the expected intensity of this dataset, or apply other laws to this dataset, which part of the module should I modify to achieve my goal?
Also, I expect you to update the code related to N-body MNIST experiments.
Thank you very much in advance.