Extrapolation in volume and custom derivatives#
There may be some cases where a custom derivative needs to be written. Most likely, this will be when you can’t write the necessary averages in terms of the observable and a specific Hamiltonian or thermodynamic conjugate variable. For instance, this might occur when extrapolating over volume in the NVT ensemble (see https://aip.scitation.org/doi/10.1063/5.0014282). In these cases, it is useful to know how to write custom derivative calculation functions and merge them back into the capabilities of the thermoextrap package.
This example will be based on the code found in thermoextrap.volume, with sections of that file reproduced here for pedagogical purposes.
First, it is useful to know the structure of how extrapolation coefficients (derivatives) are calculated in thermoextrap. Handily, there is a class called thermoextrap.models.Derivatives that uses functions or arrays of functions to compute derivatives at specific orders. Typically, these functions are generated using sympy based on known relationships between the derivatives and moments of the observable and Hamiltonian. However, they can instead be specified manually.
%matplotlib inline
import cmomy
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
rng = cmomy.random.default_rng(seed=0)
Our goal here will be to extrapolate in volume (or 1D length) for the 1D ideal gas test system.
To extrapolate in volume we need observable values at each snapshot along with values of the virial at each snapshot. For the ideal gas system, the virial is given by \(W = -\sum_{i} \frac{\mathrm{d}U}{\mathrm{d}x_{i}} x_{i} = -\sum_{i} a x_{i}\), which for \(a=1\) leads us to \(W = -N \langle x \rangle\), where the average is over all particles of a given configuration. For the purposes of extrapolation, we need a dimensionless virial, so \(W\) will be multiplied by \(\beta\).
For any observable, there is a unique term that appears in the derivative with respect to system size or volume, specifically \(\langle \sum_{i} \frac{\partial Y}{\partial x_i} x_{i}\) where \(Y\) is the observable and \(x_i\) is the \(i^\mathrm{th}\) degree of freedom. In the ideal gas system at hand, \(Y = \frac{1}{N} \sum_{i} x_{i}\) so that this correction term is equal to the observable itself, \(Y = \langle x \rangle\), with again the average here only over all particles. It is interesting that the virial is then exactly \(N\) times larger than this correction term. This is typical, as the virial will scale with the number of degrees of freedom in the system. As a result the correction term is 1000 times smaller than the other two terms appearing in the first derivative with respect to system size, \(L\). However, the other two terms nearly cancel, resulting in a small difference the same order of magnitude as the correction term. For any system you work with, you should check the magnitude of the unique correction term, either visually or by deriving and computing it.
Getting back to the code, below is a custom derivative calculation function appearing in xtrap.volume.py. This will actually work for any extrapolation in volume for the NVT ensemble, not just for the 1D ideal gas test system.
import thermoextrap as xtrap
class VolumeDerivFuncs:
"""
Calculates specific derivative values at refV with data x and W.
Only go to first order for volume extrapolation.
Here W represents the virial instead of the potential energy.
"""
def __getitem__(self, order):
# Check to make sure not going past first order
if order > 1:
raise ValueError(
"Volume derivatives cannot go past 1st order"
+ " and received %i" % order
+ "\n(because would need derivatives of forces)"
)
return self.create_deriv_func(order)
@staticmethod
def create_deriv_func(order):
def func(W, xW, dxdq, volume, ndim=1):
"""Dxdq is <sum_{i=1}^N dy/dx_i x_i>"""
# NOTE: W here has beta in it:
# that is W <- beta * virial
# Zeroth order derivative
if order == 0:
deriv_val = xW[0]
# First order derivative
else:
deriv_val = (-xW[0] * W[1] + xW[1] + dxdq) / (volume * ndim)
return deriv_val
return func
There are two key components in the above class: a __getitem__ method and a create_deriv_func method. __getitem__ is used to specify which order of derivative we want and return it - if this was an array of functions, we could just store the array and access them for the appropriate order. Here, this method just ensures that we don’t try and go past 1st order in our derivatives (because that is not implemented) and returns the derivative function shown. The create_deriv_func returns the appropriate derivative at the specified order. The function it returns takes as inputs the virial moments (would be array of potential energy moments in other scenarios), array of moments of the product of \(x\) and \(W\), an explicit derivative term for of the observable with respect to the extrapolation variable, and the volume (or \(L\) for the 1D ideal gas here). The last argument, ndim, is the dimensionality of the system, so just 1 for the 1D ideal gas, but would be 3 in more typical systems.
# Import idealgas module
from thermoextrap import idealgas
# Set temperature and volumes
beta = 1.0
volumes = np.arange(0.5, 10.0, 0.5)
# Define our reference size to extrapolated from
volume_ref = volumes[9]
npart = 1000 # Number of particles (in single configuration)
nconfig = 100_000 # Number of configurations
# And generate new data at this reference size, using a beta of 1.0 for simplicity
xdatavol, udatavol = idealgas.generate_data((nconfig, npart), beta=beta, vol=volume_ref)
# Wrap with xarray
xdatavol = xr.DataArray(xdatavol, dims=["rec"])
udatavol = xr.DataArray(udatavol, dims=["rec"])
# For the IG model, the virial W is the same as the negative of the number of particles multiplied by the average position
# (for a=1)
# And we have to multiply by beta to make the virial dimensionless
wdatavol = -1.0 * npart * xdatavol
Next we create and train our volume extrapolation model and check its outputs. You can look at factory_extrapmodel(), but it’s nearly identical to that in thermoextrap.beta, but uses specialized data callbacks and Derivatives models. We have already demonstrated above how to specify custom derivatives for the Derivatives model.
To create a special callback in a data object, follow thermoextrap.volume.VolumeDataCallback, reproduced below. Passing such callbacks to a data object through the argument meta is a very flexible way to adjust the data structures and information provided when calculating derivatives. Here, we extend the data class to also track the quantity dxdqv, or the specific derivative of the observable with respect to the extrapolation variable and is unique for every observable (see https://aip.scitation.org/doi/10.1063/5.0014282).
Basically, the argument meta takes one of these custom callbacks with the implemented methods shown below (ignore dxdqv, though, as that is specific to this case). You must implement methods to define the names of extra parameters param_names, to resample, and to pass extra information to the derivative calculation function deriv_args. Note that the method deriv_args should mirror the input structure in the above custom derivative function, with extra arguments from the callback added to what is typically passed through the deriv_args function of the base data class (specifically, the potential energy moments and \(xU\) moments, which here are virial and \(xW\) moments).
import attrs
from module_utilities import cached
@attrs.define
class VolumeDataCallback(xtrap.data.DataCallbackABC):
"""object to handle callbacks of metadata"""
volume: float = attrs.field(validator=attrs.validators.instance_of(float))
dxdqv: xr.DataArray = attrs.field(
validator=attrs.validators.instance_of(xr.DataArray)
)
ndim: int = attrs.field(default=3, validator=attrs.validators.instance_of(int))
def check(self, data) -> None:
pass
@property
def param_names(self):
return ["volume", "dxdqv", "ndim"]
@cached.meth
def dxdq(self, rec_dim, skipna):
return self.dxdqv.mean(rec_dim, skipna=skipna)
def resample(self, data, meta_kws, indices, **kws):
return self.new_like(dxdqv=self.dxdqv[indices])
def derivs_args(self, data, derivs_args):
return (
*tuple(derivs_args),
self.dxdq(rec_dim=data.rec_dim, skipna=data.skipna),
self.volume,
self.ndim,
)
xemv = xtrap.volume.factory_extrapmodel(
volume=volume_ref,
# this is actually w
uv=wdatavol,
xv=xdatavol,
# dxdqv = single observation of sum(dx/dq_i q_i) where q_i is the ith coordinate
# for ideal gas, this is just xdata
dxdqv=xdatavol,
ndim=1,
)
xemv.data.resample(nrep=100)
DataValues(meta=VolumeDataCallback(volume=5.0, dxdqv=<xarray.DataArray (rep: 100, rec: 100000)> Size: 80MB
array([[0.9643, 0.945 , 0.972 , ..., 0.9794, 1.0292, 0.9493],
[1.0295, 0.9327, 0.9943, ..., 0.9955, 0.9677, 0.968 ],
[0.961 , 0.9611, 0.9295, ..., 0.9535, 0.9424, 0.9151],
...,
[0.9817, 0.9893, 0.971 , ..., 0.999 , 0.9989, 0.989 ],
[1.0031, 0.9609, 0.9303, ..., 0.9588, 0.9442, 0.9889],
[0.9899, 0.9262, 1.0107, ..., 0.9793, 0.9503, 0.9679]])
Dimensions without coordinates: rep, rec, ndim=1), uv=<xarray.DataArray (rep: 100, rec: 100000)> Size: 80MB
array([[ -964.3265, -945.0338, -972.0076, ..., -979.4306, -1029.1505,
-949.2677],
[-1029.4963, -932.7293, -994.2952, ..., -995.4652, -967.722 ,
-968.0245],
[ -961.0341, -961.056 , -929.5435, ..., -953.4791, -942.4041,
-915.1136],
...,
[ -981.7355, -989.2521, -970.9601, ..., -998.9505, -998.9235,
-988.9737],
[-1003.0948, -960.9222, -930.2621, ..., -958.7821, -944.1744,
-988.8942],
[ -989.8773, -926.2173, -1010.7326, ..., -979.3198, -950.2932,
-967.8904]])
Dimensions without coordinates: rep, rec, xv=<xarray.DataArray (rep: 100, rec: 100000)> Size: 80MB
array([[0.9643, 0.945 , 0.972 , ..., 0.9794, 1.0292, 0.9493],
[1.0295, 0.9327, 0.9943, ..., 0.9955, 0.9677, 0.968 ],
[0.961 , 0.9611, 0.9295, ..., 0.9535, 0.9424, 0.9151],
...,
[0.9817, 0.9893, 0.971 , ..., 0.999 , 0.9989, 0.989 ],
[1.0031, 0.9609, 0.9303, ..., 0.9588, 0.9442, 0.9889],
[0.9899, 0.9262, 1.0107, ..., 0.9793, 0.9503, 0.9679]])
Dimensions without coordinates: rep, rec, order=1, rec_dim='rec', umom_dim='umom', deriv_dim=None, skipna=False, chunk=None, compute=False, build_aves_kws={}, x_is_u=False)
# Check the computation of derivatives
print("Model parameters (derivatives):\n", xemv.derivs(norm=False).values)
# Finally, look at predictions
print("Model predictions:\n", xemv.predict(volumes[:4]).values)
# And bootstrapped uncertainties
print(
"Bootstrapped uncertainties in predictions:\n",
xemv.resample(nrep=100).predict(volumes[:4]).std("rep").values,
)
Model parameters (derivatives):
[0.966 0.0269]
Model predictions:
[0.8451 0.8586 0.872 0.8854]
Bootstrapped uncertainties in predictions:
[0.0035 0.0031 0.0027 0.0023]
fig, ax = plt.subplots()
nsampvals = np.array((10.0 * np.ones(5)) ** np.arange(1, 6), dtype=int)
nsampcolors = plt.cm.viridis(np.arange(0.0, 1.0, float(1.0 / len(nsampvals))))
# First plot the analytical result
ax.plot(volumes, idealgas.x_ave(beta, volumes), "k--", linewidth=2.0)
# And the infinite sampling results for first order extrapolation
trueExtrap, trueDerivs = idealgas.x_vol_extrap(1, volume_ref, volumes, beta=beta)
ax.plot(volumes, trueExtrap, "r-", zorder=0)
print(f"True extrapolation coefficients: {trueDerivs}")
for i, n in enumerate(nsampvals):
thisinds = rng.choice(len(xdatavol), size=n, replace=False)
# Get parameters for extrapolation model with this data by training it - the parameters are the derivatives
xemv = xtrap.volume.factory_extrapmodel(
volume=volume_ref,
# this is actually w
uv=wdatavol[thisinds],
xv=xdatavol[thisinds],
dxdqv=xdatavol[thisinds],
ndim=1,
)
out = xemv.predict(volumes)
print("\t With N_configs = %6i: %s" % (n, str(xemv.derivs(norm=False).values)))
out.plot(marker="s", ms=4, color=nsampcolors[i], ls="None", label=f"N={n}", ax=ax)
ax.set_ylabel(r"$\langle x \rangle$")
ax.set_xlabel(r"$L$")
ax.annotate("O(1) Extrapolation", xy=(0.4, 0.7), xycoords="axes fraction", fontsize=9)
ax.set_ylim((-1.0, 4.0))
ax.yaxis.set_major_locator(mpl.ticker.MaxNLocator(nbins=4, prune="both"))
fig.tight_layout()
fig.subplots_adjust(hspace=0.0)
ax.set_title(None)
plt.legend()
plt.show()
True extrapolation coefficients: [0.9661 0.0274]
With N_configs = 10: [0.9572 0.0866]
With N_configs = 100: [0.9637 0.0683]
With N_configs = 1000: [0.9654 0.0307]
With N_configs = 10000: [0.9662 0.0315]
With N_configs = 100000: [0.966 0.0269]
