def init_fn(key, R, box, mass=f32(1.0), **kwargs):
N, dim = R.shape
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
V = jnp.sqrt(_kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - jnp.mean(V * mass, axis=0, keepdims=True) / mass
KE = quantity.kinetic_energy(V, mass)
# The box position is defined via pos = (1 / d) log V / V_0.
zero = jnp.zeros((), dtype=R.dtype)
one = jnp.ones((), dtype=R.dtype)
box_position = zero
box_velocity = zero
box_mass = dim * (N + 1) * kT * barostat_kwargs['tau']**2 * one
KE_box = quantity.kinetic_energy(box_velocity, box_mass)
if jnp.isscalar(box) or box.ndim == 0:
# TODO(schsam): This is necessary because of JAX issue #5849.
box = jnp.eye(R.shape[-1]) * box
return NPTNoseHooverState(R, V, force_fn(R, box=box, **kwargs), mass,
box, box_position, box_velocity, box_mass,
barostat.initialize(1, KE_box, _kT),
thermostat.initialize(R.size, KE, _kT)) # pytype: disable=wrong-arg-count
def init_fun(key: Array,
R: Array,
velocity_scale: float = f32(1.0),
mass=f32(1.0),
**kwargs) -> NVEState:
V = np.sqrt(velocity_scale) * random.normal(
key, R.shape, dtype=R.dtype)
mass = quantity.canonicalize_mass(mass)
return NVEState(R, V, force(R, **kwargs) / mass, mass) # pytype: disable=wrong-arg-count
def init_fn(key, R, mass=f32(1), **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
key, split = random.split(key)
V = np.sqrt(_kT / mass) * random.normal(split, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
return NVTLangevinState(R, V, force_fn(R, **kwargs), mass, key) # pytype: disable=wrong-arg-count
def init_fn(key, R, mass=f32(1), T_initial=f32(1)):
mass = quantity.canonicalize_mass(mass)
key, split = random.split(key)
V = np.sqrt(T_initial / mass) * random.normal(
split, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
return NVTLangevinState(R, V, force_fn(R, t=f32(0)), mass, key)
def init_fn(key, R, mass=f32(1.0), **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
V = jnp.sqrt(_kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - jnp.mean(V * mass, axis=0, keepdims=True) / mass
KE = quantity.kinetic_energy(V, mass)
return NVTNoseHooverState(R, V, force_fn(R, **kwargs), mass,
chain_fns.initialize(R.size, KE, _kT)) # pytype: disable=wrong-arg-count
def init_fn(key, R, mass=f32(1), T_initial=None, **kwargs):
if T_initial is None:
T_initial = T_schedule(0.0)
mass = quantity.canonicalize_mass(mass)
key, split = random.split(key)
V = np.sqrt(T_initial / mass) * random.normal(
split, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
return NVTLangevinState(R, V, force_fn(R, t=f32(0), **kwargs), mass,
key) # pytype: disable=wrong-arg-count
def init_fun(key, R, mass=f32(1.0), T_initial=f32(1.0)):
mass = quantity.canonicalize_mass(mass)
V = np.sqrt(T_initial / mass) * random.normal(
key, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
# Nose-Hoover parameters.
xi = np.zeros(chain_length, R.dtype)
v_xi = np.zeros(chain_length, R.dtype)
DOF, = static_cast(R.shape[0] * R.shape[1])
Q = T_initial * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
return NVTNoseHooverState(R, V, mass, KE, xi, v_xi, Q)
def init_fn(key, R, mass=f32(1.0), **kwargs):
_kT = kT if 'kT' not in kwargs else kwargs['kT']
mass = quantity.canonicalize_mass(mass)
V = np.sqrt(_kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
# Nose-Hoover parameters.
xi = np.zeros(chain_length, R.dtype)
v_xi = np.zeros(chain_length, R.dtype)
# TODO(schsam): Really, it seems like Q should be set by the goal
# temperature rather than the initial temperature.
DOF = f32(R.shape[0] * R.shape[1])
Q = _kT * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
F = force_fn(R, **kwargs)
return NVTNoseHooverState(R, V, F, mass, KE, xi, v_xi, Q) # pytype: disable=wrong-arg-count
def init_fun(key, R, mass=f32(1.0), T_initial=None):
if T_initial is None:
T_initial = T_schedule(0.0)
mass = quantity.canonicalize_mass(mass)
V = np.sqrt(T_initial / mass) * random.normal(
key, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
KE = quantity.kinetic_energy(V, mass)
# Nose-Hoover parameters.
xi = np.zeros(chain_length, R.dtype)
v_xi = np.zeros(chain_length, R.dtype)
# TODO(schsam): Really, it seems like Q should be set by the goal
# temperature rather than the initial temperature.
DOF, = static_cast(R.shape[0] * R.shape[1])
Q = T_initial * tau**f32(2) * np.ones(chain_length, dtype=R.dtype)
Q = ops.index_update(Q, 0, Q[0] * DOF)
return NVTNoseHooverState(R, V, mass, KE, xi, v_xi, Q) # pytype: disable=wrong-arg-count
def init_fn(key, R, mass=f32(1)):
mass = quantity.canonicalize_mass(mass)
return BrownianState(R, mass, key) # pytype: disable=wrong-arg-count
def init_fn(key, R, mass=f32(1)):
mass = quantity.canonicalize_mass(mass)
return BrownianState(R, mass, key)
def init_fun(key, R, velocity_scale=f32(1.0), mass=f32(1.0)):
V = np.sqrt(velocity_scale) * random.normal(
key, R.shape, dtype=R.dtype)
mass = quantity.canonicalize_mass(mass)
return NVEState(R, V, force(R) / mass, mass)
def test_canonicalize_mass(self):
assert quantity.canonicalize_mass(3.0) == 3.0
assert quantity.canonicalize_mass(f32(3.0)) == f32(3.0)
assert quantity.canonicalize_mass(f64(3.0)) == f64(3.0)
def init_fn(key, R, kT, mass=f32(1.0), **kwargs):
mass = quantity.canonicalize_mass(mass)
V = jnp.sqrt(kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - jnp.mean(V * mass, axis=0, keepdims=True) / mass
return NVEState(R, V, force_fn(R, **kwargs), mass) # pytype: disable=wrong-arg-count
def init_fn(key, R, kT, mass=f32(1.0), **kwargs):
mass = quantity.canonicalize_mass(mass)
V = np.sqrt(kT / mass) * random.normal(key, R.shape, dtype=R.dtype)
V = V - np.mean(V, axis=0, keepdims=True)
return NVEState(R, V, force_fn(R, **kwargs), mass)