From 35a23acc1c744e49964cd591e62be376ab704c79 Mon Sep 17 00:00:00 2001 From: Logan Ward Date: Wed, 13 Dec 2023 15:28:27 -0500 Subject: [PATCH] Add function which returns the MVN (#37) --- jitterbug/model/base.py | 46 +++++++++++++++++++---------- jitterbug/model/linear.py | 15 +++------- jitterbug/model/linear_internals.py | 12 ++++++++ tests/models/test_linear.py | 9 ++++++ 4 files changed, 56 insertions(+), 26 deletions(-) diff --git a/jitterbug/model/base.py b/jitterbug/model/base.py index db8a0df..61581c8 100644 --- a/jitterbug/model/base.py +++ b/jitterbug/model/base.py @@ -5,6 +5,7 @@ from pathlib import Path import os +import ase import numpy as np from scipy import stats from ase import Atoms @@ -15,6 +16,9 @@ class EnergyModel: """Base class for functions which predict energy given molecular structure""" + reference: ase.Atoms + """Structure for which we will be computing Hessians""" + def train(self, data: list[Atoms]) -> object: """Produce an energy model given observations of energies @@ -43,6 +47,31 @@ def sample_hessians(self, model: object, num_samples: int) -> list[np.ndarray]: Returns: A list of 2D hessians """ + n_params = len(self.reference) * 3 + dist = self.get_hessian_distribution(model) + diag_ind = np.diag_indices(n_params) + triu_ind = np.triu_indices(n_params) + output = [] + for sample in dist.rvs(size=num_samples): + # Fill in a 2D version + sample_hessian = np.zeros((n_params, n_params)) + sample_hessian[triu_ind] = sample + + # Make it symmetric + sample_hessian += sample_hessian.T + sample_hessian[diag_ind] /= 2 + + output.append(sample_hessian) + return output + + def get_hessian_distribution(self, model: object) -> stats.multivariate_normal: + """Get a multi-variate normal distribution of the independent parameters for the hessian + + Args: + model: Model trained by this class + Returns: + A MVN distribution for the upper triangle of the Hessian matrix (in row-major order) + """ raise NotImplementedError() @@ -105,7 +134,7 @@ def mean_hessian(self, models: list[Calculator]) -> np.ndarray: # Return the mean return np.mean(hessians, axis=0) - def sample_hessians(self, models: list[Calculator], num_samples: int) -> list[np.ndarray]: + def get_hessian_distribution(self, models: list[Calculator]) -> list[np.ndarray]: # Run all calculators hessians = [self.compute_hessian(self.reference, calc) for calc in models] @@ -116,17 +145,4 @@ def sample_hessians(self, models: list[Calculator], num_samples: int) -> list[np hessian_covar = np.cov(hessians_flat, rowvar=False) # Generate samples - hessian_mvn = stats.multivariate_normal(hessian_mean, hessian_covar, allow_singular=True) - diag_indices = np.diag_indices(hessians[0].shape[0]) - output = [] - for sample in hessian_mvn.rvs(size=num_samples): - # Fill in a 2D version - sample_hessian = np.zeros_like(hessians[0]) - sample_hessian[triu_ind] = sample - - # Make it symmetric - sample_hessian += sample_hessian.T - sample_hessian[diag_indices] /= 2 - - output.append(sample_hessian) - return output + return stats.multivariate_normal(hessian_mean, hessian_covar, allow_singular=True) diff --git a/jitterbug/model/linear.py b/jitterbug/model/linear.py index c418927..5b77b9d 100644 --- a/jitterbug/model/linear.py +++ b/jitterbug/model/linear.py @@ -6,6 +6,7 @@ """ import numpy as np from ase import Atoms +from scipy.stats import multivariate_normal from sklearn.linear_model import ARDRegression from sklearn.linear_model._base import LinearModel @@ -83,7 +84,7 @@ def train(self, data: list[Atoms]) -> LinearModel: def mean_hessian(self, model: LinearModel) -> np.ndarray: return self._params_to_hessian(model.coef_) - def sample_hessians(self, model: LinearModel, num_samples: int) -> list[np.ndarray]: + def get_hessian_distribution(self, model: LinearModel) -> multivariate_normal: # Get the covariance matrix if not hasattr(model, 'sigma_'): # pragma: no-coverage raise ValueError(f'Sampling only possible with Bayesian regressors. You trained a {type(model)}') @@ -101,14 +102,8 @@ def sample_hessians(self, model: LinearModel, num_samples: int) -> list[np.ndarr sigma = model.sigma_ # Sample the model parameters - params = np.random.multivariate_normal(model.coef_, sigma, size=num_samples) - - # Assemble them into Hessians - output = [] - for param in params: - hessian = self._params_to_hessian(param) - output.append(hessian) - return output + n_coords = len(self.reference) * 3 + return multivariate_normal(model.coef_[n_coords:], sigma[n_coords:, n_coords:], allow_singular=True) def _params_to_hessian(self, param: np.ndarray) -> np.ndarray: """Convert the parameters for the linear model into a Hessian @@ -128,6 +123,4 @@ def _params_to_hessian(self, param: np.ndarray) -> np.ndarray: hessian[triu_inds] = param[n_coords:] # The first n_coords terms are the linear part hessian[off_diag_triu_inds] /= 2 hessian.T[triu_inds] = hessian[triu_inds] - # v = np.sqrt(self.reference.get_masses()).repeat(3).reshape(-1, 1) - # hessian /= np.dot(v, v.T) return hessian diff --git a/jitterbug/model/linear_internals.py b/jitterbug/model/linear_internals.py index 9d94cfe..080f59d 100644 --- a/jitterbug/model/linear_internals.py +++ b/jitterbug/model/linear_internals.py @@ -9,6 +9,7 @@ from ase import io as aseio from geometric.molecule import Molecule from geometric.internal import DelocalizedInternalCoordinates as DIC +from scipy import stats from sklearn.linear_model import ARDRegression from sklearn.linear_model._base import LinearModel from .base import EnergyModel @@ -175,6 +176,17 @@ def sample_hessians(self, model: LinearModel, num_samples: int) -> list[np.ndarr output.append(hessian) return output + def get_hessian_distribution(self, model: LinearModel) -> stats.multivariate_normal: + # I'll go back and figure out the correct math for transforming the covariance matrix of the internal hessian + # to that of the cartesian Hessian. We'll just refit + ind = np.triu_indices(len(self.reference) * 3) + samples = [x[ind] for x in self.sample_hessians(model, 128)] + return stats.multivariate_normal( + np.mean(samples, axis=0), + np.cov(samples, rowvar=False), + allow_singular=True + ) + def _params_to_hessian(self, param: np.ndarray) -> np.ndarray: """Convert the parameters for the linear model into a Hessian diff --git a/tests/models/test_linear.py b/tests/models/test_linear.py index 36fed22..c674319 100644 --- a/tests/models/test_linear.py +++ b/tests/models/test_linear.py @@ -49,6 +49,15 @@ def test_linear_model(train_set, model_type, num_params): assert len(hessians) == 32 assert np.isclose(hessians[0], hessians[0].T).all() + # Make sure the underlying distribution has the same mean + dist = model.get_hessian_distribution(hessian_model) + ind = np.triu_indices(9) + assert np.isclose( + np.mean(hessians, axis=0)[ind], + dist.mean, + atol=10 # Does not have to be close. Our sampling size is pretty small + ).all() + # Only test accuracy with IC harmonic. Other one's trash if isinstance(model, ICHarmonicModel): vib_data = VibrationsData.from_2d(reference, hessians[0])