Skip to content

Commit

Permalink
Add linear terms to harmonic model (#6)
Browse files Browse the repository at this point in the history 
* Refactor and add linear terms

Fixes #2
Fixes #4

* Make the second-order terms 1/2 the first order

* Update with the newer model

* Update the proof of concept notebooks
  • Loading branch information
@WardLT
WardLT authored Aug 30, 2023
1 parent 7a31d68 commit 93f77b3
Show file tree
Hide file tree
Showing 4 changed files with 252 additions and 169 deletions.
85 changes: 60 additions & 25 deletions jitterbug/model/linear.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,34 +12,48 @@
from .base import EnergyModel


def get_displacement_matrix(atoms: Atoms, reference: Atoms) -> np.ndarray:
"""Get the displacements of a structure from a reference
in the order, as used in a Hessian calculation.
def get_model_inputs(atoms: Atoms, reference: Atoms) -> np.ndarray:
"""Get the inputs for the model, which are derived from the displacements
of the structure with respect to a reference.
Args:
atoms: Displaced structure
reference: Reference structure
Returns:
Vector of displacements
Vector of displacements in the same order as the
"""

# Compute the displacements
disp_matrix = (reference.positions - atoms.positions).flatten()
disp_matrix = disp_matrix[:, None] * disp_matrix[None, :]
# Compute the displacements and the products of displacement
disp_matrix = (atoms.positions - reference.positions).flatten()
disp_prod_matrix = disp_matrix[:, None] * disp_matrix[None, :]

# Multiply the off-axis terms by two, as they appear twice in the energy model
n_terms = len(atoms) * 3
off_diag = np.triu_indices(n_terms, k=1)
disp_matrix[off_diag] *= 2
disp_prod_matrix[off_diag] *= 2

# Return the upper triangular matrix
return disp_matrix[np.triu_indices(n_terms)]
# Append the displacements and products of displacements
return np.concatenate([
disp_matrix,
disp_prod_matrix[np.triu_indices(n_terms)] / 2
], axis=0)


class LinearHessianModel(EnergyModel):
"""Fits a model for energy using linear regression
class HarmonicModel(EnergyModel):
"""Expresses energy as a Harmonic model (i.e., 2nd degree Taylor series)
Implicitly treats all elements of the Hessian matrix as unrelated
Contains a total of :math:`3N + 3N(3N+1)/2` terms in total, where :math:`N`
is the number of atoms in the molecule. The first :math:`3N` correspond to the
linear terms of the model, which are known as the Jacobian matrix, and the
latter are from the quadratic terms, which are half of the symmetric Hessian matrix.
Implicitly treats all terms of the model as unrelated, which is the worst case
for trying to fit the energy of a molecule. However, it is still possible to fit
the model with a reduced number of terms if we assume that most terms are near zero.
The energy model is:
:math:`E = E_0 + \\sum_i J_i \\delta_i + \\frac{1}{2}\\sum_{i,j} H_{i,j}\\delta_i\\delta_j`
Args:
reference: Fully-relaxed structure used as the reference
Expand All @@ -50,10 +64,9 @@ def __init__(self, reference: Atoms, regressor: type[LinearModel] = ARDRegressio
self.reference = reference
self.regressor = regressor

def train(self, data: list[Atoms]) -> ARDRegression:
def train(self, data: list[Atoms]) -> LinearModel:
# X: Displacement vectors for each
x = [get_displacement_matrix(atoms, self.reference) for atoms in data]
x = np.multiply(x, 0.5)
x = [get_model_inputs(atoms, self.reference) for atoms in data]

# Y: Subtract off the reference energy
ref_energy = self.reference.get_potential_energy()
Expand All @@ -68,12 +81,28 @@ def train(self, data: list[Atoms]) -> ARDRegression:

return model

def mean_hessian(self, model: ARDRegression) -> np.ndarray:
def mean_hessian(self, model: LinearModel) -> np.ndarray:
return self._params_to_hessian(model.coef_)

def sample_hessians(self, model: ARDRegression, num_samples: int) -> list[np.ndarray]:
def sample_hessians(self, model: LinearModel, num_samples: int) -> list[np.ndarray]:
# 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)}')
if isinstance(model, ARDRegression):
# The sigma matrix may be zero for high-precision terms
n_terms = len(model.coef_)
nonzero_terms = model.lambda_ < model.threshold_lambda

# Replace those terms (Thanks: https://stackoverflow.com/a/73176327/2593278)
sigma = np.zeros((n_terms, n_terms))
sub_sigma = sigma[nonzero_terms, :]
sub_sigma[:, nonzero_terms] = model.sigma_
sigma[nonzero_terms, :] = sub_sigma
else:
sigma = model.sigma_

# Sample the model parameters
params = np.random.multivariate_normal(model.coef_, model.sigma_, size=num_samples)
params = np.random.multivariate_normal(model.coef_, sigma, size=num_samples)

# Assemble them into Hessians
output = []
Expand All @@ -83,15 +112,21 @@ def sample_hessians(self, model: ARDRegression, num_samples: int) -> list[np.nda
return output

def _params_to_hessian(self, param: np.ndarray) -> np.ndarray:
"""Convert the parameters for the linear model into a Hessian"""
"""Convert the parameters for the linear model into a Hessian
Args:
param: Coefficients of the linear model
Returns:
The harmonic terms expressed as a Hessian matrix
"""
# Get the parameters
n_terms = len(self.reference) * 3
triu_inds = np.triu_indices(n_terms)
off_diag_triu_inds = np.triu_indices(n_terms, k=1)
n_coords = len(self.reference) * 3
triu_inds = np.triu_indices(n_coords)
off_diag_triu_inds = np.triu_indices(n_coords, k=1)

# Assemble the hessian
hessian = np.zeros((n_terms, n_terms))
hessian[triu_inds] = param
hessian = np.zeros((n_coords, n_coords))
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]
return hessian
107 changes: 28 additions & 79 deletions notebooks/proof-of-concept/1_compute-random-offsets.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"id": "c6a28419-6831-4197-8973-00c5591e19cb",
"metadata": {
"tags": []
Expand All @@ -38,7 +38,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"id": "c6be56c5-a460-4acd-9b89-8c3d9c812f5f",
"metadata": {
"tags": [
Expand All @@ -64,7 +64,7 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"id": "0b6794cd-477f-45a1-b96f-2332804ddb20",
"metadata": {},
"outputs": [],
Expand All @@ -83,23 +83,12 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"id": "ad9fd725-b1ba-4fec-ae41-959be0e540b3",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"Atoms(symbols='O2N4C8H10', pbc=False, forces=..., calculator=SinglePointCalculator(...))"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"outputs": [],
"source": [
"atoms = read(Path('data') / 'exact' / f'{run_name}.xyz')\n",
"atoms"
Expand All @@ -111,7 +100,7 @@
"metadata": {},
"source": [
"## Compute many random energies\n",
"Compute $3N(3N+1)/2 + 1$ energies with displacements sampled [on the unit sphere](https://mathoverflow.net/questions/24688/efficiently-sampling-points-uniformly-from-the-surface-of-an-n-sphere). This is enough to fit the Hessian exactly plus a little more"
"Compute $3N + 3N(3N+1)/2 + 1$ energies with displacements sampled [on the unit sphere](https://mathoverflow.net/questions/24688/efficiently-sampling-points-uniformly-from-the-surface-of-an-n-sphere). This is enough to fit the Jacobian and Hessian exactly plus a little more"
]
},
{
Expand All @@ -124,7 +113,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"id": "23502eea-0974-4248-8f19-e85447069c61",
"metadata": {
"tags": []
Expand All @@ -137,7 +126,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"id": "bf1366fc-d9a7-4a98-b9c9-cb3a0209b406",
"metadata": {
"tags": []
Expand All @@ -157,7 +146,7 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"id": "d4f21e81-5ec3-4877-a4d1-402077be2ee8",
"metadata": {
"tags": []
Expand All @@ -179,22 +168,12 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": null,
"id": "0915595d-133a-43df-84fc-4ff6a3b538ea",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Memory set to 3.815 GiB by Python driver.\n",
" Threads set to 12 by Python driver.\n"
]
}
],
"outputs": [],
"source": [
"calc = Psi4(method=method, basis=basis, num_threads=threads, memory='4096MB')"
]
Expand All @@ -209,42 +188,26 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": null,
"id": "e2a28593-2634-4bb7-ae5b-8f557937bda1",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Need to run 2629 calculations for full accuracy.\n"
]
}
],
"outputs": [],
"source": [
"n_atoms = len(atoms)\n",
"to_compute = 3 * n_atoms * (3 * n_atoms + 1) // 2 + 1\n",
"to_compute = 3 * n_atoms + 3 * n_atoms * (3 * n_atoms + 1) // 2 + 1\n",
"print(f'Need to run {to_compute} calculations for full accuracy.')"
]
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": null,
"id": "8bf40523-dcaa-4046-a9c6-74e35178e87f",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Already done 1. 2628 left to do.\n"
]
}
],
"outputs": [],
"source": [
"with connect(db_path) as db:\n",
" done = len(db)\n",
Expand All @@ -256,41 +219,19 @@
"execution_count": null,
"id": "a6fa1b33-defc-4b35-895d-052eb64453fb",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
" 0%| | 0/2629 [00:00<?, ?it/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
" Memory set to 3.815 GiB by Python driver.\n",
" Threads set to 12 by Python driver.\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" 10%|███████████████████████ | 252/2629 [28:25<4:10:48, 6.33s/it]"
]
}
],
"outputs": [],
"source": [
"pbar = tqdm(total=to_compute)\n",
"pbar.update(done)\n",
"for i in range(to_compute - done):\n",
" # Sample a perturbation\n",
" disp = np.random.normal(0, 1, size=(n_atoms * 3))\n",
" disp -= disp.mean()\n",
" disp /= np.linalg.norm(disp)\n",
" disp *= step_size\n",
" disp = disp.reshape((-1, 3))\n",
" \n",
" # Subtract off any translation\n",
" disp -= disp.mean(axis=0)[None, :]\n",
"\n",
" # Make the new atoms\n",
" new_atoms = atoms.copy()\n",
Expand All @@ -304,6 +245,14 @@
"\n",
" pbar.update(1)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "785add47-39b5-4d7e-9d92-0375c8128171",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
Expand Down
Loading

0 comments on commit 93f77b3

Please sign in to comment.