Langevin adaptative corrector #109
Conversation
| @@ -1,3 +1,5 @@ | |||
| from diffusion_for_multi_scale_molecular_dynamics.generators.adaptative_corrector import \ | |||
| AdaptativeCorrectorGenerator | |||
There was a problem hiding this comment.
"Adaptative" -> "Adaptive"
| NoiseParameters | ||
|
|
||
|
|
||
| class AdaptativeCorrectorGenerator(LangevinGenerator): |
There was a problem hiding this comment.
"Adaptatif" is french. "Adaptive" is a better choice in English.
There was a problem hiding this comment.
I changed it. But interesting grammatical fact: both adaptative and adaptive are in the oxford dictionary and act as synonyms.
| @@ -518,7 +522,8 @@ def corrector_step( | |||
| ) -> AXL: | |||
| """Corrector Step. | |||
|
|
|||
| Note this is not affecting the atom types. Only the reduced coordinates and lattice vectors. | |||
| Note this dones not affect the atom types unless specified with the atom_type_transition_in_corrector argument. | |||
| ) | ||
| self.corrector_r = noise_parameters.corrector_r | ||
|
|
||
| def predictor_step( |
There was a problem hiding this comment.
If I'm reading this right, this predictor_step is identical to the predictor_step in LangevinGenerator, with the only difference that x_im1 = x_i. I note that the two methods have diverged (you copy-pasted an older version I believe).
I propose that you create methods
LangevinGenerator._relative_coordinates_update_predictor_step
andLangevinGenerator._relative_coordinates_update_corrector_step
In LangevinGenerator, both methods could just be dummies that call LangevinGenerator._relative_coordinates_update.
In the class AdaptativeCorrectorGenerator, you could overload _relative_coordinates_update_predictor_step to return its own input.
This way, we avoid a lot of duplicated code.
There was a problem hiding this comment.
I agree. I changed the code accordingly.
|
|
||
| return composition_im1 | ||
|
|
||
| def corrector_step( |
There was a problem hiding this comment.
This method also seems like it is very close to LangevinGenerator.corrector_step.
I propose that you create a method LangevinGenerator._get_corrector_step_epsilon. In LangevinGenerator, it could just return the correct entry from the tabulated values, and you could overload this method here to compute epsilon in a more sophisticated way.
We could get rid of the array sqrt_2_epsilon in LangevinDynamics: surely it's not taking a sqrt that slows us down.
Again, this would avoid a lot of duplication.
There was a problem hiding this comment.
changed. I did not delete sqrt_2_epsilon yet. It requires checking a few tests to make sure it is not called anywhere. We could either do it right now or leave it as a TODO as part of a code cleanup effort.
| ) | ||
|
|
||
| # to compute epsilon_i, we need the norm of the score. We average over the atoms. | ||
| relative_coordinates_sigma_score_norm = ( |
There was a problem hiding this comment.
The norm in the paper is calculated over all pixels.
I think you are doing this:
and I think the right way of doing it is this:
The paper talks about averaging over the batch dimension; you are averaging over the atoms within a single sample.
I don't think we should do this just yet...
There was a problem hiding this comment.
I wasn't sure what was the best approach to take the norm, so I went with something reasonable. Let's stick to the paper recommendations for now.
| from tests.generators.test_langevin_generator import TestLangevinGenerator | ||
|
|
||
|
|
||
| class TestAdaptativeCorrectorGenerator(TestLangevinGenerator): |
There was a problem hiding this comment.
Adaptative -> Adaptive
Also, the word "adaptative" still apears in the main code at the following locations:
./src/diffusion_for_multi_scale_molecular_dynamics/noise_schedulers/noise_parameters.py: # Step size scaling for the Adaptative Corrector Generator. Default value comes from github implementation
./src/diffusion_for_multi_scale_molecular_dynamics/generators/adaptive_corrector.py: """Langevin Dynamics Generator using only a corrector step with adaptative step size for relative coordinates.
./src/diffusion_for_multi_scale_molecular_dynamics/generators/sde_position_generator.py: adaptative: bool = False
./src/diffusion_for_multi_scale_molecular_dynamics/generators/sde_position_generator.py: adaptive=self.sampling_parameters.adaptative,
| z: torch.Tensor, | ||
| ) -> torch.Tensor: | ||
| """Compute the size of the corrector step for the relative coordinates update.""" | ||
| del sigma_i # sigma_i is not needed to compute \epsilon_i because it is pre-tabulated |
There was a problem hiding this comment.
Why delete these variables? It's ok not to use them.
| @@ -518,7 +583,8 @@ def corrector_step( | |||
| ) -> AXL: | |||
| """Corrector Step. | |||
|
|
|||
| Note this is not affecting the atom types. Only the reduced coordinates and lattice vectors. | |||
| Note this does not not affect the atom types unless specified with the atom_type_transition_in_corrector | |||
This implements the Langevin Corrector generator from
SCORE-BASED GENERATIVE MODELING THROUGH STOCHASTIC DIFFERENTIAL EQUATIONS
algorithm 4 in appendix G. I renamed it Adaptative Corrector to minimize name collisions.
The differences from the Langevin Generator only apply to the coordinates sampler, not the atom type.