Loss

crystal_diffusion/analysis/exploding_variance_analysis.py

@@ -8,38 +8,29 @@
 
 from crystal_diffusion import ANALYSIS_RESULTS_DIR
 from crystal_diffusion.analysis import PLEASANT_FIG_SIZE, PLOT_STYLE_PATH
-from crystal_diffusion.samplers.time_sampler import TimeParameters, TimeSampler
 from crystal_diffusion.samplers.variance_sampler import (
-    ExplodingVarianceSampler, VarianceParameters)
+    ExplodingVarianceSampler, NoiseParameters)
 from crystal_diffusion.score.wrapped_gaussian_score import \
     get_sigma_normalized_score
 
 plt.style.use(PLOT_STYLE_PATH)
 
 if __name__ == '__main__':
 
-    variance_parameters = VarianceParameters()
-    time_parameters = TimeParameters(total_time_steps=1000)
+    noise_parameters = NoiseParameters(total_time_steps=1000)
+    variance_sampler = ExplodingVarianceSampler(noise_parameters=noise_parameters)
 
-    time_sampler = TimeSampler(time_parameters=time_parameters)
-    variance_sampler = ExplodingVarianceSampler(variance_parameters=variance_parameters,
-                                                time_sampler=time_sampler)
+    noise = variance_sampler.get_all_noise()
 
-    indices = torch.arange(time_parameters.total_time_steps)
-    times = time_sampler.get_time_steps(indices)
-    sigmas = torch.sqrt(variance_sampler.get_variances(indices))
-    gs = torch.sqrt(variance_sampler.get_g_squared_factors(indices[1:]))
-
-    # A first figure to compare the "smart" and the "brute force" results
     fig1 = plt.figure(figsize=PLEASANT_FIG_SIZE)
     fig1.suptitle("Noise Schedule")
 
     ax1 = fig1.add_subplot(221)
     ax2 = fig1.add_subplot(223)
     ax3 = fig1.add_subplot(122)
 
-    ax1.plot(times, sigmas, '-', c='k', lw=2)
-    ax2.plot(times[1:], gs, '-', c='k', lw=2)
+    ax1.plot(noise.time, noise.sigma, '-', c='k', lw=2)
+    ax2.plot(noise.time[1:], noise.g[1:], '-', c='k', lw=2)
 
     ax1.set_ylabel('$\\sigma(t)$')
     ax2.set_ylabel('$g(t)$')
@@ -55,17 +46,18 @@
 
     kmax = 4
     indices = torch.tensor([1, 250, 750, 999])
-    times = time_sampler.get_time_steps(indices)
-    sigmas = torch.sqrt(variance_sampler.get_variances(indices))
-    gs_squared = variance_sampler.get_g_squared_factors(indices)
+
+    times = noise.time.take(indices)
+    sigmas = noise.sigma.take(indices)
+    gs_squared = noise.g_squared.take(indices)
 
     for t, sigma in zip(times, sigmas):
-        target_scores = get_sigma_normalized_score(relative_positions,
-                                                   torch.ones_like(relative_positions) * sigma,
-                                                   kmax=kmax)
-        ax3.plot(relative_positions, sigma * target_scores, label=f"t = {t:3.2f}")
+        target_sigma_normalized_scores = get_sigma_normalized_score(relative_positions,
+                                                                    torch.ones_like(relative_positions) * sigma,
+                                                                    kmax=kmax)
+        ax3.plot(relative_positions, target_sigma_normalized_scores, label=f"t = {t:3.2f}")
 
-    ax3.set_title("Target Noise")
+    ax3.set_title("Target Normalized Score")
     ax3.set_xlabel("relative position, u")
     ax3.set_ylabel("$\\sigma(t) \\times S(u, t)$")
     ax3.legend(loc=0)

crystal_diffusion/analysis/target_score_analysis.py

@@ -10,8 +10,8 @@
 from crystal_diffusion import ANALYSIS_RESULTS_DIR
 from crystal_diffusion.analysis import PLEASANT_FIG_SIZE, PLOT_STYLE_PATH
 from crystal_diffusion.score.wrapped_gaussian_score import (
-    SIGMA_THRESHOLD, get_expected_sigma_normalized_score_brute_force,
-    get_sigma_normalized_score)
+    SIGMA_THRESHOLD, get_sigma_normalized_score,
+    get_sigma_normalized_score_brute_force)
 
 plt.style.use(PLOT_STYLE_PATH)
 
@@ -32,7 +32,7 @@
         sigma = sigma_factor * SIGMA_THRESHOLD
 
         sigmas = torch.ones_like(relative_positions) * sigma
-        list_scores_brute = np.array([get_expected_sigma_normalized_score_brute_force(u, sigma) for u in list_u])
+        list_scores_brute = np.array([get_sigma_normalized_score_brute_force(u, sigma) for u in list_u])
         list_scores = get_sigma_normalized_score(relative_positions, sigmas, kmax=kmax).numpy()
         error = list_scores - list_scores_brute
 
@@ -73,7 +73,7 @@
                      ms=ms, c=color, lw=2, alpha=0.25, label=f'kmax = {kmax}')
 
         list_scores_brute = np.array([
-            get_expected_sigma_normalized_score_brute_force(u, sigma, kmax=4 * kmax) for sigma in sigmas])
+            get_sigma_normalized_score_brute_force(u, sigma, kmax=4 * kmax) for sigma in sigmas])
         ax4.semilogy(sigma_factors, list_scores_brute, 'o-',
                      ms=ms, c=color, lw=2, alpha=0.25, label=f'kmax = {4 * kmax}')
 

crystal_diffusion/models/optimizer.py

@@ -0,0 +1,41 @@
+import logging
+from dataclasses import dataclass
+from enum import Enum
+
+import torch
+from torch import optim
+
+logger = logging.getLogger(__name__)
+
+
+class ValidOptimizerNames(Enum):
+    """Valid optimizer names."""
+    adam = "adam"
+    sgd = "sgd"
+
+
+@dataclass(kw_only=True)
+class OptimizerParameters:
+    """Parameters for the optimizer."""
+    name: ValidOptimizerNames
+    learning_rate: float
+
+
+def load_optimizer(hyper_params: OptimizerParameters, model: torch.nn.Module) -> optim.Optimizer:
+    """Instantiate the optimizer.
+
+    Args:
+        hyper_params : hyperparameters defining the optimizer
+        model : A neural network model.
+
+    Returns:
+        optimizer : The optimizer for the given model
+    """
+    match hyper_params.name:
+        case ValidOptimizerNames.adam:
+            optimizer = optim.Adam(model.parameters(), lr=hyper_params.learning_rate)
+        case ValidOptimizerNames.sgd:
+            optimizer = optim.SGD(model.parameters(), lr=hyper_params.learning_rate)
+        case _:
+            raise ValueError(f"optimizer {hyper_params.name} not supported")
+    return optimizer