From 0798f299cc8a25f680702b53e571157a29b09bb2 Mon Sep 17 00:00:00 2001 From: Logan Ward Date: Thu, 5 Oct 2023 12:27:05 -0400 Subject: [PATCH] Add scaling factors to calculator, simplify interfaces --- jitterbug/model/soap.py | 41 ++++++++++++++++++++++++++------------- tests/models/test_soap.py | 4 ++-- 2 files changed, 30 insertions(+), 15 deletions(-) diff --git a/jitterbug/model/soap.py b/jitterbug/model/soap.py index b2d7b7a..fdf76ac 100644 --- a/jitterbug/model/soap.py +++ b/jitterbug/model/soap.py @@ -15,12 +15,20 @@ class InducedKernelGPR(ExactGP): """Gaussian process regression model with an induced kernel Predicts the energy for each atom as a function of its descriptors + + Args: + batch_x: Example input batch of descriptors + inducing_x: Starting points for the reference points of the kernel + likelihood: Likelihood function used to describe the noise + use_ard: Whether to employ a different length scale parameter for each descriptor, + a technique known as Automatic Relevance Detection (ARD) """ - def __init__(self, batch_x: torch.Tensor, batch_y: torch.Tensor, inducing_x: torch.Tensor, likelihood: GaussianLikelihood, use_adr: bool): + def __init__(self, batch_x: torch.Tensor, inducing_x: torch.Tensor, likelihood: GaussianLikelihood, use_ard: bool): + batch_y = torch.zeros(batch_x.shape[0], dtype=batch_x.dtype) super().__init__(batch_x, batch_y, likelihood) self.mean_module = ConstantMean() - self.base_covar_module = ScaleKernel(RBFKernel(has_lengthscale=True, ard_num_dims=batch_x.shape[-1] if use_adr else None)) + self.base_covar_module = ScaleKernel(RBFKernel(has_lengthscale=True, ard_num_dims=batch_x.shape[-1] if use_ard else None)) self.covar_module = InducingPointKernel(self.base_covar_module, inducing_points=inducing_x.clone(), likelihood=likelihood) def forward(self, x) -> MultivariateNormal: @@ -29,13 +37,13 @@ def forward(self, x) -> MultivariateNormal: return MultivariateNormal(mean_x, covar_x) -def make_gpr_model(train_descriptors: np.ndarray, num_inducing_points: int, use_adr_kernel: bool = False) -> InducedKernelGPR: +def make_gpr_model(train_descriptors: np.ndarray, num_inducing_points: int, use_ard_kernel: bool = False) -> InducedKernelGPR: """Make the GPR model for a certain atomic system Args: train_descriptors: 3D array of all training points (num_configurations, num_atoms, num_descriptors) num_inducing_points: Number of inducing points to use in the kernel. More points, more complex model - use_adr_kernel: Whether to use a different length scale parameter for each descriptor + use_ard_kernel: Whether to use a different length scale parameter for each descriptor Returns: Model which can predict energy given descriptors for a single configuration """ @@ -46,13 +54,11 @@ def make_gpr_model(train_descriptors: np.ndarray, num_inducing_points: int, use_ inducing_points = descriptors[inducing_inds, :] # Make the model - # TODO (wardlt): Consider making a batch size of more than one. Do we do that by just supplying all of the training data - batch_y = torch.zeros(descriptors.shape[0]) # Use a mean of zero for the training points + # TODO (wardlt): Consider making a batch size of more than one. We currently use the entire training set in a single batch return InducedKernelGPR( batch_x=torch.from_numpy(descriptors), - batch_y=batch_y, # Returns a scalar for each point inducing_x=torch.from_numpy(inducing_points), - use_adr=use_adr_kernel, + use_ard=use_ard_kernel, likelihood=GaussianLikelihood() ) @@ -67,7 +73,7 @@ def train_model(model: InducedKernelGPR, train_x: np.ndarray, train_y: np.ndarra steps: Number of interactions over all training points learning_rate: Learning rate used for the optimizer Returns: - Mean loss over all iterations + Mean loss over each iteration """ # Convert the data to numpy @@ -112,12 +118,20 @@ def train_model(model: InducedKernelGPR, train_x: np.ndarray, train_y: np.ndarra class SOAPCalculator(Calculator): - """Calculator which uses a GPR model trained using SOAP descriptors""" + """Calculator which uses a GPR model trained using SOAP descriptors + + Args: + model (InducedKernelGPR): A machine learning model which takes descriptors as inputs + soap (SOAP): Tool used to compute the descriptors + scaling (tuple[float, float]): A offset and factor with which to adjust the energy per atom predictions, + which are typically he mean and standard deviation of energy per atom across the training set. + """ implemented_properties = ['energy', 'forces', 'energies'] default_parameters = { 'model': None, - 'soap': None + 'soap': None, + 'scaling': (0., 1.) } def calculate(self, atoms=None, properties=('energy', 'forces', 'energies'), @@ -129,11 +143,12 @@ def calculate(self, atoms=None, properties=('energy', 'forces', 'energies'), d_desc_d_pos = torch.from_numpy(d_desc_d_pos) # Run inference + offset, scale = self.parameters['scaling'] model: InducedKernelGPR = self.parameters['model'] model.eval() # Ensure we're in eval mode pred_energies_dist = model(desc) pred_energies = pred_energies_dist.mean - pred_energy = torch.sum(pred_energies) + pred_energy = torch.sum(pred_energies) * scale + offset # Compute the forces # See: https://singroup.github.io/dscribe/latest/tutorials/machine_learning/forces_and_energies.html @@ -146,7 +161,7 @@ def calculate(self, atoms=None, properties=('energy', 'forces', 'energies'), grad_outputs=torch.ones_like(pred_energy), )[0] # Derivative of the energy with respect to the descriptors for each center d_energy_d_center_d_pos = torch.einsum('ijkl,il->ijk', d_desc_d_pos, d_energy_d_desc) # Derivative for each center with respect to each atom - pred_forces = -d_energy_d_center_d_pos.sum(dim=0) # Total effect on each center from each atom + pred_forces = -d_energy_d_center_d_pos.sum(dim=0) * scale # Total effect on each center from each atom # Store the results self.results['forces'] = pred_forces.detach().numpy() diff --git a/tests/models/test_soap.py b/tests/models/test_soap.py index 95dd3b3..ba61c94 100644 --- a/tests/models/test_soap.py +++ b/tests/models/test_soap.py @@ -23,8 +23,8 @@ def descriptors(train_set, soap): @mark.parametrize('use_adr', [True, False]) -def test_make_model(use_adr, descriptors): - model = make_gpr_model(descriptors, 4, use_adr_kernel=use_adr) +def test_make_model(use_adr, descriptors, train_set): + model = make_gpr_model(descriptors, 4, use_ard_kernel=use_adr) # Evaluate on a single point model.eval()