IMLCV.base package

IMLCV.base package#

Submodules#

IMLCV.base.CV module#

class IMLCV.base.CV.CV(*, cv, mapped=False, atomic=False, _combine_dims=None, _stack_dims=None)#

Bases: MyPyTreeNode

atomic: bool = False#

batch()#

Return type:: CV

property batch_dim#

property batched#

static combine(*cvs, flatten=False)#

merges a list of CVs into a single CV. The dimenisions are stored such that it can later be split into separated CVs. The CVs are combined over the last dimension

Return type:: CV

property combine_dims#

static create(cv, mapped=False, atomic=False, combine_dims=None, stack_dims=None)#

cv: Array#

property dim#

mapped: bool = False#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

property shape#

property size#

split(flatten=False)#

inverse operation of combine

Return type:: list[CV]

static stack(*cvs)#

stacks a list of CVs into a single CV. The dimenisions are stored such that it can later be unstacked into separated CVs. The CVs are stacked over the batch dimension

Return type:: CV

property stack_dims#

unbatch()#

Return type:: CV

unstack()#

Return type:: list[CV]

class IMLCV.base.CV.CollectiveVariable(*, f, metric, jac=<function jacrev>)#

Bases: MyPyTreeNode

compute_cv(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

f: CvTrans#

jac(argnums=0, has_aux=False, holomorphic=False, allow_int=False)#

Jacobian of fun evaluated row-by-row using reverse-mode AD.

Parameters:

fun (Callable) – Function whose Jacobian is to be computed.
argnums (int | Sequence[int]) – Optional, integer or sequence of integers. Specifies which positional argument(s) to differentiate with respect to (default 0).
has_aux (bool) – Optional, bool. Indicates whether fun returns a pair where the first element is considered the output of the mathematical function to be differentiated and the second element is auxiliary data. Default False.
holomorphic (bool) – Optional, bool. Indicates whether fun is promised to be holomorphic. Default False.
allow_int (bool) – Optional, bool. Whether to allow differentiating with respect to integer valued inputs. The gradient of an integer input will have a trivial vector-space dtype (float0). Default False.

Return type:

Callable

Returns:

A function with the same arguments as fun, that evaluates the Jacobian of fun using reverse-mode automatic differentiation. If has_aux is True then a pair of (jacobian, auxiliary_data) is returned.

>>> import jax
>>> import jax.numpy as jnp
>>>
>>> def f(x):
...   return jnp.asarray(
...     [x[0], 5*x[2], 4*x[1]**2 - 2*x[2], x[2] * jnp.sin(x[0])])
...
>>> print(jax.jacrev(f)(jnp.array([1., 2., 3.])))
[[ 1.       0.       0.     ]
 [ 0.       0.       5.     ]
 [ 0.      16.      -2.     ]
 [ 1.6209   0.       0.84147]]

static load(file, **kwargs)#

Return type:: CollectiveVariable

metric: CvMetric#

property n#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

save(file)#

class IMLCV.base.CV.CvFun(*, kwargs=<factory>, static_kwargs=<factory>, forward=None, backward=None)#

Bases: CvFunBase

backward: Optional[Callable[[CV, NeighbourList | None, CV | None], CV]] = None#

forward: Optional[Callable[[CV, NeighbourList | None, CV | None], CV]] = None#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

class IMLCV.base.CV.CvFunBase(*, kwargs=<factory>, static_kwargs=<factory>)#

Bases: ABC, MyPyTreeNode

compute_cv(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

kwargs: dict#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

static_kwargs: dict#

class IMLCV.base.CV.CvMetric(*, bounding_box, periodicities)#

Bases: MyPyTreeNode

class to keep track of topology of given CV. Identifies the periodicitie of CVs and maps to unit square with correct peridicities

bounding_box: Array#

static bounds_from_cv(cv_0, percentile=0.1, weights=None, rho=None, margin=None, chunk_size=None, n=400, macro_chunk=5000, verbose=True)#

classmethod create(periodicities=None, bounding_box=None)#

Return type:: CvMetric

difference(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

static get_n(samples_per_bin, samples, n_dims, max_bins=None, max_bins_per_dim=30)#

grid(n=30, bounds=None, margin=0.1, indexing='ij')#

forms regular grid in mapped space. If coordinate is periodic, last rows are ommited.

Parameters:

n – number of points in each dim
map – boolean. True: work in mapped space (default), False: calculate grid in space without metric
endpoints – if

Returns:

meshgrid and vector with distances between points

map(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

min_cv(cv)#

property ndim#

norm(x1, x2, k=1.0)#

periodic_wrap(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

periodicities: Array#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

unmap(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

class IMLCV.base.CV.CvTrans(*, trans)#

Bases: MyPyTreeNode

compute_cv(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

static from_cv_function(f, static_argnames=None, check_input=True, **kwargs)#

Return type:: CvTrans

replace(**updates)#: Returns a new object replacing the specified fields with new values.

trans: _ParralelCvTrans | _SerialCvTrans#

class IMLCV.base.CV.NeighbourList(*, sp_orig, info, update, atom_indices=None, op_cell=None, op_coor=None, op_center=None, ijk_indices=None, _padding_bools=None)#

Bases: MyPyTreeNode

apply_fun_neighbour(sp, func, r_cut=None, fill_value=0, reduce='full', exclude_self=False, chunk_size_neigbourgs=None, chunk_size_atoms=None, chunk_size_batch=None, shmap=False, split_z=False, shmap_kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0),), mesh=Mesh(device_ids=array([0]), axis_names=('i',), axis_types=(Auto,))))#

apply_fun_neighbour_pair(sp, func_double, func_single=<function NeighbourList.<lambda>>, r_cut=None, fill_value=0.0, reduce='full', split_z=False, exclude_self=True, unique=True, chunk_size_neigbourgs=None, chunk_size_atoms=None, chunk_size_batch=None, shmap=False, shmap_kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0), ), mesh=Mesh(device_ids=array([0]), axis_names=('i', ), axis_types=(Auto, ))))#

Return type:: tuple[Array, TypeVar(S)]

Args:#

func_single(r_ij, atom_index_j) = .. func_double( p_ij, atom_index_j, data_j, p_ik, atom_index_k, data_k) = …

atom_indices: Array | None = None#

batch()#

property batch_dim#

property batched#

canonicalized_sp(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

static create(r_cut, sp_orig, atom_indices=None, r_skin=None, ijk_indices=None, z_array=None, nxyz=None, op_cell=None, op_coor=None, op_center=None, stack_dims=None, num_neighs=None)#

ijk_indices: Array | None = None#

info: NeighbourListInfo#

property needs_calculation#

needs_update(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

neighbour_pos(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

nl_split_z(p)#

nneighs(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

op_cell: Array | None = None#

op_center: Array | None = None#

op_coor: Array | None = None#

property padding_bools#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

property shape#

sp_orig: SystemParams | None#

static stack(*nls)#

Return type:: NeighbourList

property stack_dims#

unstack()#

Return type:: list[NeighbourList]

update: NeighbourListUpdate#

update_nl(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

class IMLCV.base.CV.NeighbourListInfo(*, r_cut, r_skin, z_array=None, z_unique=None, num_z_unique=None)#

Bases: MyPyTreeNode

static create(r_cut, z_array, r_skin=None)#

nl_split_z(p)#

Return type:: tuple[Array, list[Array], list[TypeVar(T)]]

num_z_unique: tuple[int] | None = None#

r_cut: float#

r_skin: float#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

z_array: tuple[int] | None = None#

z_unique: tuple[int] | None = None#

class IMLCV.base.CV.NeighbourListUpdate(*, nxyz=None, stack_dims=None, num_neighs=None)#

Bases: MyPyTreeNode

static create(nxyz=None, stack_dims=None, num_neighs=None)#

num_neighs: int | None = None#

nxyz: tuple[int] | None = None#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

stack_dims: tuple[int] | None = None#

class IMLCV.base.CV.ShmapKwargs(*, axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=None, mesh=None)#

Bases: MyPyTreeNode

axis: int = 0#

axis_name: str | None = 'i'#

static create(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=None, mesh=None)#

device_get: bool = True#

devices: tuple[Any] = None#

explicit_shmap: bool = False#

mesh: Mesh | None = None#

n_devices: int | None = None#

out_axes: int = 0#

pmap: bool = False#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

verbose: bool = False#

class IMLCV.base.CV.SystemParams(*, coordinates, cell=None)#

Bases: MyPyTreeNode

angles(deg=True)#

Return type:: Array

apply_canonicalize(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

apply_minkowski_reduction(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

apply_rotation(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

apply_wrap(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

batch()#

Return type:: SystemParams

property batch_dim#

property batched#

canonicalize(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

cell: Array | None = None#

coordinates: Array#

get_neighbour_list(info, chunk_size=None, chunk_size_inner=100, verbose=False, shmap=False, shmap_kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0),), mesh=Mesh(device_ids=array([0]), axis_names=('i',), axis_types=(Auto,))), only_update=False)#

Return type:: NeighbourList | None

min_distance(index_1, index_2)#

minkowski_reduce(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

replace(**updates)#: Returns a new object replacing the specified fields with new values.

rotate_cell(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

property shape#

static stack(*sps)#

Return type:: SystemParams

super_cell(n, info=None)#

Return type:: tuple[SystemParams, NeighbourListInfo | None]

to_absolute(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

to_ase(static_trajectory_info)#

to_relative(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

unbatch()#

Return type:: SystemParams

volume()#

wrap_positions(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

IMLCV.base.CV.macro_chunk_map(f, y, nl=None, ft=None, y_t=None, nl_t=None, macro_chunk=1000, verbose=False, print_every=10, jit_f=True)#

Return type:: list[TypeVar(X2, CV, SystemParams, NeighbourList)] | tuple[list[TypeVar(X2, CV, SystemParams, NeighbourList)], list[TypeVar(X2, CV, SystemParams, NeighbourList)] | None]

IMLCV.base.CV.macro_chunk_map_fun(f, y, w=None, nl=None, ft=None, y_t=None, nl_t=None, macro_chunk=1000, verbose=False, chunk_func=None, chunk_func_init_args=None, w_t=None, print_every=10, jit_f=True)#

Return type:: TypeVar(T)

IMLCV.base.CV.padded_shard_map(f, kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0),), mesh=Mesh(device_ids=array([0]), axis_names=('i',), axis_types=(Auto,))))#

IMLCV.base.CV.padded_vmap(f, chunk_size=None, axis=0, out_axes=0, vmap=True, verbose=False)#

IMLCV.base.CVDiscovery module#

class IMLCV.base.CVDiscovery.CombineTransformer(*, outdim, descriptor=None, pre_scale=True, post_scale=True, transformers)#

Bases: Transformer

static create(transformers)#

Return type:: CombineTransformer

replace(**updates)#: Returns a new object replacing the specified fields with new values.

transformers: list[Transformer]#

class IMLCV.base.CVDiscovery.CvTransTransformer(*, outdim, descriptor=None, pre_scale=True, post_scale=True, trans)#

Bases: Transformer

replace(**updates)#: Returns a new object replacing the specified fields with new values.

trans: CvTrans#

class IMLCV.base.CVDiscovery.IdentityTransformer(*, outdim, descriptor=None, pre_scale=True, post_scale=True)#

Bases: Transformer

replace(**updates)#: Returns a new object replacing the specified fields with new values.

class IMLCV.base.CVDiscovery.Transformer(*, outdim, descriptor=None, pre_scale=True, post_scale=True)#

Bases: MyPyTreeNode

descriptor: CvTrans | None = None#

fit(dlo, chunk_size=None, plot=True, plot_folder=None, shmap=True, percentile=5.0, jac=<function jacrev>, transform_FES=True, koopman=True, max_fes_bias=None, n_max=100000.0, samples_per_bin=20, min_samples_per_bin=3, verbose=True, cv_titles=True, vmax=0.03808799176039228, macro_chunk=1000, shmap_kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0), ), mesh=Mesh(device_ids=array([0]), axis_names=('i', ), axis_types=(Auto, ))), **kwargs)#

Return type:: tuple[list[CV], CollectiveVariable, Bias]

outdim: int#

static plot_app(collective_variables, cv_data=None, biases=None, indicate_plots='lightblue', duplicate_cv_data=True, name=None, labels=None, cv_titles=True, data_titles=None, color_trajectories=False, margin=0.1, plot_FES=False, T=None, vmin=0, vmax=0.03808799176039228, dpi=300, n_max_bias=1000000.0, row_color=None, macro_chunk=10000, cmap='jet', offset=True, bar_label='FES [kJ/mol]')#: Plot the app for the CV discovery. all 1d and 2d plots are plotted directly, 3d or higher are plotted as 2d slices.

post_scale: bool = True#

pre_fit(dlo, chunk_size=None, shmap=True, shmap_kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0),), mesh=Mesh(device_ids=array([0]), axis_names=('i',), axis_types=(Auto,))), verbose=False, macro_chunk=10000)#

pre_scale: bool = True#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

static static_fit(transformer, **kwargs)#

IMLCV.base.MdEngine module#

MD engine class peforms MD simulations in a given NVT/NPT ensemble.

class IMLCV.base.MdEngine.MDEngine(bias, energy, sp, static_trajectory_info, trajectory_info=None, trajectory_file=None, time0=<factory>, step=1, nl=None)#

Bases: ABC

Base class for MD engine.

bias: Bias#

classmethod create(bias, energy, static_trajectory_info, trajectory_info=None, trajectory_file=None, sp=None, **kwargs)#

Return type:: Self

energy: Energy#

get_bias(sp, gpos=False, vtens=False, shmap=False, use_jac=False, push_jac=False, rel=False, shmap_kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0),), mesh=Mesh(device_ids=array([0]), axis_names=('i',), axis_types=(Auto,))))#

Return type:: tuple[CV, EnergyResult]

get_energy(sp, gpos=False, vtens=False, manual_vir=None)#

Return type:: EnergyResult

get_trajectory()#

Return type:: TrajectoryInfo

static load(file, **kwargs)#

Return type:: MDEngine

new_bias(bias, **kwargs)#

Return type:: MDEngine

nl: NeighbourList | None = None#

r_skin = 1.889726133921252#

run(steps)#

run the integrator for a given number of steps.

Parameters:: steps – number of MD steps

save(file)#

save_step(T=None, P=None, t=None, err=None, cv=None, e_bias=None, e_pot=None, gpos_rmsd=None, gpos_bias_rmsd=None, canonicalize=False)#

sp: SystemParams#

static_trajectory_info: StaticMdInfo#

step: int = 1#

time0: float#

trajectory_file: Path | None = None#

trajectory_info: TrajectoryInfo | None = None#

update_nl()#

class IMLCV.base.MdEngine.StaticMdInfo(timestep, T, timecon_thermo, atomic_numbers, r_cut=None, P=None, timecon_baro=None, write_step=100, equilibration=None, screen_log=1000, invalid=False, max_grad=0.04031059429903557)#

Bases: object

P: float | None = None#

T: float#

atomic_numbers: Array#

property barostat#

equilibration: float | None = None#

invalid: bool = False#

static load(filename)#

Return type:: StaticMdInfo

property masses#

max_grad: float | None = 0.04031059429903557#

property neighbour_list_info: NeighbourListInfo#

r_cut: float | None = None#

save(filename)#

screen_log: int = 1000#

property thermostat#

timecon_baro: float | None = None#

timecon_thermo: float#

timestep: float#

write_step: int = 100#

class IMLCV.base.MdEngine.TrajectoryInfo(*, _positions, _cell=None, _charges=None, _e_pot=None, _e_bias=None, _w=None, _rho=None, _cv=None, _cv_orig=None, _T=None, _P=None, _err=None, _t=None, _capacity=-1, _size=-1, _finished=False, _invalid=False)#

Bases: MyPyTreeNode

property CV: CV | None#

property CV_orig: CV | None#

property P: Array | None#

property T: Array | None#

property cell: Array | None#

property charges: Array | None#

static create(positions, cell=None, charges=None, e_pot=None, e_bias=None, cv=None, cv_orig=None, w=None, rho=None, T=None, P=None, err=None, t=None, capacity=-1, size=-1, finished=False, invalid=False)#

Return type:: TrajectoryInfo

property cv: Array | None#

property cv_orig: Array | None#

property e_bias: Array | None#

property e_pot: Array | None#

property err: Array | None#

static load(filename)#

Return type:: TrajectoryInfo

property positions: Array | None#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

property rho: Array | None#

save(filename)#

property shape#

property sp: SystemParams#

property t: Array | None#

property volume#

property w: Array | None#

IMLCV.base.Observable module#

class IMLCV.base.Observable.Observable(*, cv_round, rounds, rnd, common_bias, collective_variable, time_per_bin=82682.74667329364)#

Bases: object

class to convert data and CVs to different thermodynamic/ kinetic observables.

collective_variable: CollectiveVariable#

common_bias: Bias#

static create(rounds, rnd=None, cv_round=None, collective_variable=None)#

Return type:: Observable

cv_round: int#

fes_bias(plot=True, fes=None, max_bias=None, choice='rbf', num_rnds=8, start_r=1, rbf_kernel='gaussian', rbf_degree=None, samples_per_bin=5, min_samples_per_bin=1, chunk_size=None, macro_chunk=10000, update_bounding_box=True, n_max=100000.0, min_traj_length=None, margin=0.1, only_finished=True, shmap=False, thermolib=False, lag_n=30, out=100000, vmax=0.03808799176039228, koopman=True, verbose=True, koopman_wham=None, executors=Executors.training, direct_bias=False)#

fes_nd_thermolib(num_rnds=4, start_r=1, update_bounding_box=True, samples_per_bin=5, min_samples_per_bin=1, chunk_size=None, n_max=100000.0, n=None, min_traj_length=None, dlo=None, directory=None, temp=None, shmap=False, only_finished=True, bounds_percentile=1, max_bias=None, rbf_kernel='gaussian', rbf_degree=None)#

fes_nd_weights(num_rnds=4, out=100000, lag_n=10, start_r=1, min_traj_length=None, only_finished=True, chunk_size=None, macro_chunk=1000, n_max=100000.0, cv_round=None, directory=None, koopman=True, verbose=True, max_bias=0.03808799176039228, kooopman_wham=None, samples_per_bin=5, min_samples_per_bin=1, executors=Executors.training, direct_bias=False)#

rnd: int#

rounds: Rounds#

time_per_bin: float = 82682.74667329364#

IMLCV.base.UnitsConstants module#

IMLCV.base.bias module#

class IMLCV.base.bias.Bias(*, collective_variable, start=0, step=1, finalized=False, slice_exponent=1.0, log_exp_slice=True, slice_mean=False)#

Bases: ABC, MyPyTreeNode

base class for biased MD runs.

apply(cvs, shmap=False, macro_chunk_size=10000)#

bounds_from_bias(T, sign=1.0, margin=1e-10, n=50)#

collective_variable: CollectiveVariable#

compute_from_cv(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

compute_from_system_params(**kwargs: ~typing.~P) → T#

Return type:: TypeVar(T)

static create(*args, **kwargs)#

finalized: bool = False#

kl_divergence(other, T, symmetric=True, margin=0.2, sign=1.0, n=100)#

static load(filename)#

Return type:: Bias

log_exp_slice: bool = True#

plot(name=None, traj=None, dlo_kwargs=None, dlo=None, vmax=0.03808799176039228, map=False, inverted=False, margin=0.1, dpi=300, T=300.0, **kwargs)#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

resample(cv_grid=None, n=40, margin=0.3)#

Return type:: Bias

save(filename, cv_file=None)#

slice(T, inverted=True, vmax=None, n_max_bias=100000.0, margin=0.2, macro_chunk=10000, offset=True)#

Return type:: dict[int, dict[tuple[int], Bias]]

slice_exponent: float = 1.0#

slice_mean: bool = False#

start: int | None = 0#

step: int | None = 1#

update_bias(md)#

update the bias.

Can only change the properties from _get_args

Return type:: Bias

exception IMLCV.base.bias.BiasError#: Bases: Exception

class IMLCV.base.bias.BiasF(*, collective_variable, start=0, step=1, finalized=False, slice_exponent=1.0, log_exp_slice=True, slice_mean=False, g=<function _constant>, static_kwargs=<factory>, kwargs=<factory>)#

Bases: Bias

Bias according to CV.

static create(cvs, g=<function _constant>, kwargs={}, static_kwargs={})#

Return type:: BiasF

g(val=0.0)#

kwargs: dict#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

static_kwargs: dict#

class IMLCV.base.bias.BiasModify(*, collective_variable, start=0, step=1, finalized=False, slice_exponent=1.0, log_exp_slice=True, slice_mean=False, fun, bias, kwargs=<factory>, static_kwargs=<factory>)#

Bases: Bias

Bias according to CV.

bias: Bias#

classmethod create(fun, bias, kwargs={}, static_kwargs={})#

Return type:: BiasModify

fun: Callable#

kwargs: dict#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

static_kwargs: dict#

update_bias(md)#

update the bias.

Can only change the properties from _get_args

Return type:: Bias

class IMLCV.base.bias.CompositeBias(*, collective_variable, start=0, step=1, finalized=False, slice_exponent=1.0, log_exp_slice=True, slice_mean=False, biases, fun)#

Bases: Bias

Class that combines several biases in one single bias.

biases: list[Bias]#

classmethod create(biases, fun=<function sum>)#

Return type:: CompositeBias

fun: Callable#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

update_bias(md)#

update the bias.

Can only change the properties from _get_args

Return type:: Bias

class IMLCV.base.bias.Energy#

Bases: object

abstract property cell: Array | None#

compute_from_system_params(sp, gpos=False, vir=False, nl=None, manual_vir=None, shmap=False, shmap_kwarg=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0),), mesh=Mesh(device_ids=array([0]), axis_names=('i',), axis_types=(Auto,))))#

Return type:: EnergyResult

abstract property coordinates: Array | None#

external_callback = True#

get_vtens_finite_difference(sp, nl, eps=1e-05, gpos=False)#

static load(filename)#

Return type:: Energy

manual_vtens = False#

property nl: NeighbourList | None#

save(filename)#

property sp: SystemParams | None#

exception IMLCV.base.bias.EnergyError#: Bases: Exception

class IMLCV.base.bias.EnergyFn(*, f, _sp=None, _nl=None, kwargs=<factory>, static_kwargs=<factory>)#

Bases: Energy, MyPyTreeNode

property cell: Array | None#

property coordinates: Array | None#

external_callback = False#

f: Callable#

kwargs: dict#

property nl#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

property sp: SystemParams | None#

static_kwargs: dict#

class IMLCV.base.bias.EnergyResult(*, energy, gpos=None, vtens=None)#

Bases: MyPyTreeNode

energy: Array#

gpos: Array | None = None#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

vtens: Array | None = None#

class IMLCV.base.bias.NoneBias(*, collective_variable, start=0, step=1, finalized=False, slice_exponent=1.0, log_exp_slice=True, slice_mean=False, g=<function _constant>, static_kwargs=<factory>, kwargs=<factory>)#

Bases: BiasF

static create(collective_variable)#

Return type:: NoneBias

replace(**updates)#: Returns a new object replacing the specified fields with new values.

class IMLCV.base.bias.PlumedEnerg#: Bases: Energy

IMLCV.base.datastructures module#

class IMLCV.base.datastructures.MyPyTreeNode(*args, **kwargs)#

Bases: object

Base class for dataclasses that should act like a JAX pytree node.

replace(**overrides)#

Return type:: TypeVar(TNode, bound= MyPyTreeNode)

IMLCV.base.datastructures.Partial_decorator(f, *partial_args, **partial_kwargs)#

IMLCV.base.datastructures.custom_jvp_decorator(f, nondiff_argnums=())#

Return type:: Callable[[ParamSpec(P)], TypeVar(T)]

IMLCV.base.datastructures.jit_decorator(f, static_argnums=None, static_argnames=None)#

IMLCV.base.datastructures.vmap_decorator(f, in_axes=0, out_axes=0)#

IMLCV.base.rounds module#

class IMLCV.base.rounds.Covariances(*, rho_00, rho_01, rho_10, rho_11, pi_s_0, pi_s_1, sigma_0, sigma_1, W_0=None, W_1=None, only_diag=False, trans_f=None, trans_g=None, symmetric=False)#

Bases: MyPyTreeNode

property C00#

property C01#

property C10#

property C11#

W_0: Array | None = None#

W_1: Array | None = None#

static create(cv_0, cv_1=None, nl=None, nl_t=None, w=None, w_t=None, calc_pi=False, macro_chunk=1000, chunk_size=None, only_diag=False, trans_f=None, trans_g=None, symmetric=False, calc_C00=True, calc_C01=True, calc_C10=True, calc_C11=True, shmap_kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0),), mesh=Mesh(device_ids=array([0]), axis_names=('i',), axis_types=(Auto,))), verbose=True)#

Return type:: Covariances

decompose(out_dim=None, sparse=True, out_eps=None)#

mask(eps_pre, max_features=2000, auto_cov_threshold=None)#

only_diag: bool = False#

property pi_0#

property pi_1#

pi_s_0: Array | None#

pi_s_1: Array | None#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

rho_00: Array#

rho_01: Array | None#

rho_10: Array | None#

rho_11: Array | None#

sigma_0: Array#

property sigma_0_inv#

sigma_1: Array | None#

property sigma_1_inv#

symmetric: bool = False#

symmetrize()#

trans_f: CvTrans | None = None#

trans_g: CvTrans | None = None#

whiten_rho(choice, epsilon=0.0001, max_features=None, verbose=False)#

Return type:: Array

class IMLCV.base.rounds.DataLoaderOutput(*, sp, cv, sti, ti, collective_variable, labels=None, nl=None, sp_t=None, nl_t=None, cv_t=None, ti_t=None, time_series=False, tau=None, bias=None, ground_bias=None, _weights=None, _weights_t=None, _rho=None, _rho_t=None, scaled_tau=False)#

Bases: MyPyTreeNode

apply_bias(bias, chunk_size=None, macro_chunk=1000, verbose=False, shmap=False)#

Return type:: list[Array]

static apply_cv(f, x, x_t=None, nl=None, nl_t=None, chunk_size=None, macro_chunk=1000, shmap=True, shmap_kwargs=ShmapKwargs(axis=0, out_axes=0, axis_name='i', n_devices=None, pmap=False, explicit_shmap=False, verbose=False, device_get=True, devices=(CpuDevice(id=0),), mesh=Mesh(device_ids=array([0]), axis_names=('i',), axis_types=(Auto,))), verbose=False, print_every=10, jit_f=True)#

Return type:: tuple[list[CV], list[CV] | None]

bias: list[Bias] | None = None#

calc_neighbours(r_cut, chunk_size=None, macro_chunk=1000, verbose=False, only_update=False, chunk_size_inner=10, max=(2, 2, 2))#

static check_w(w_stacked)#

Return type:: Array

collective_variable: CollectiveVariable#

cv: list[CV]#

cv_t: list[CV] | None = None#

filter_nans(x=None, x_t=None, macro_chunk=1000)#

get_bincount(cv_0=None, w=None, n_hist=None, n_max=100000.0, margin=0.1, chunk_size=None, use_w=False, output_labels=False, samples_per_bin=20, min_samples_per_bin=1, macro_chunk=1000)#

get_fes_bias_from_weights(cv=None, weights=None, rho=None, samples_per_bin=100, min_samples_per_bin=5, n_grid=None, n_max=1000.0, max_bias=None, chunk_size=None, macro_chunk=1000, max_bias_margin=0.2, rbf_bias=True, kernel='gaussian', collective_variable=None, set_outer_border=True, rbf_degree=None)#

static get_histo(data_nums, weights=None, log_w=False, macro_chunk=320, verbose=False, shape_mask=None, nn=-1, f_func=None)#

Return type:: Array

get_point(point, sigma=0.01, key=Array([0, 0], dtype=uint32))#

Return type:: tuple[CV, SystemParams, NeighbourList | NeighbourListInfo | None, TrajectoryInfo]

get_transformed_fes(new_cv, new_colvar, samples_per_bin=5, min_samples_per_bin=1, chunk_size=1, smoothing=0.0, max_bias=None, shmap=False, n_grid_old=50, n_grid_new=30)#

Return type:: RbfBias

ground_bias: Bias | None = None#

koopman_model(cv_0=None, cv_t=None, nl=None, nl_t=None, only_return_weights=False, symmetric=False, rho=None, w=None, rho_t=None, w_t=None, eps=1e-12, eps_pre=1e-12, max_features=5000, max_features_pre=5000, out_dim=-1, add_1=True, chunk_size=None, macro_chunk=1000, verbose=False, trans=None, T_scale=1.0, only_diag=False, calc_pi=False, scaled_tau=None, sparse=True, correlation=True, auto_cov_threshold=None)#

Return type:: KoopmanModel

koopman_weight(w=None, w_t=None, samples_per_bin=50, max_bins=100000.0, out_dim=-1, chunk_size=None, indicator_CV=True, koopman_eps=1e-05, koopman_eps_pre=0, cv_0=None, cv_t=None, macro_chunk=1000, verbose=False, max_features_koopman=5000, margin=0.1, add_1=None, only_diag=False, calc_pi=False, sparse=False, output_w_corr=False, correlation=True, return_km=False, labels=None, koopman_kwargs={})#

labels: list[Array] | None = None#

nl: list[NeighbourList] | NeighbourList | None = None#

nl_t: list[NeighbourList] | NeighbourList | None = None#

static norm_w(w_stacked, ret_n=False)#

Return type:: Array | tuple[Array, Array]

recalc(chunk_size=None, macro_chunk=1000, shmap=False, verbose=False)#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

scaled_tau: bool = False#

sp: list[SystemParams]#

sp_t: list[SystemParams] | None = None#

sti: StaticMdInfo#

tau: float | None = None#

ti: list[TrajectoryInfo]#

ti_t: list[TrajectoryInfo] | None = None#

time_series: bool = False#

transform_FES(trans, T=None, max_bias=0.03808799176039228, n_grid=25)#

wham_weight(samples_per_bin=10, n_max=100000.0, chunk_size=None, wham_eps=1e-07, cv_0=None, cv_t=None, macro_chunk=1000, verbose=False, margin=0.1, bias_cutoff=0.11426397528117684, min_f_i=1e-30, log_sum_exp=True, return_bias=False, return_std_bias=False, min_samples=3, lagrangian=False, max_sigma=0.0019043995880196139, sparse_inverse=True, inverse_sigma_weighting=False, output_bincount=True, output_free_energy=True, output_time_scaling=True, smooth_bias=False, max_bias_margin=0.1, alpha_dirichlet=1.0)#

Return type:: WeightOutput

class IMLCV.base.rounds.KoopmanModel(*, s, cov, W0, W1, argmask, shape, cv_0, cv_t, nl, nl_t, w=None, rho=None, w_t=None, rho_t=None, eps=1e-10, eps_pre=None, only_diag=False, calc_pi=True, scaled_tau=False, add_1=False, max_features=5000, max_features_pre=5000, out_dim=None, correlation_whiten=True, verbose=True, tau=None, T_scale=1.0, trans=None, constant_threshold=1e-10)#

Bases: MyPyTreeNode

T_scale: float = 1.0#

W0: Array#

W1: Array#

add_1: bool = False#

argmask: Array | None#

calc_pi: bool = True#

constant_threshold: float = 1e-10#

correlation_whiten: bool = True#

cov: Covariances#

static create(w, rho, w_t, rho_t, cv_0, cv_t, nl=None, nl_t=None, add_1=True, eps=1e-14, eps_pre=0, symmetric=False, out_dim=-1, max_features=5000, max_features_pre=5000, tau=None, macro_chunk=1000, chunk_size=None, verbose=True, trans=None, T_scale=1.0, only_diag=False, calc_pi=False, use_scipy=False, auto_cov_threshold=None, sparse=True, scaled_tau=False, only_return_weights=False, correlation_whiten=False, out_eps=None, constant_threshold=1e-10)#

cv_0: list[CV] | list[SystemParams]#

cv_t: list[CV] | list[SystemParams]#

eps: float = 1e-10#

eps_pre: float | None = None#

f(out_dim=None, remove_constant=True, n_skip=None)#

g(out_dim=None, n_skip=None, remove_constant=True)#

get_n_skip(n_skip=None, remove_constant=True)#

koopman_weight(verbose=False, chunk_size=None, macro_chunk=1000, retarget=True, epsilon=1e-06, max_entropy=True, out_dim=None)#

Return type:: tuple[list[Array], list[Array], list[Array] | None, bool]

max_features: int = 5000#

max_features_pre: int = 5000#

nl: list[NeighbourList] | NeighbourList | None#

nl_t: list[NeighbourList] | NeighbourList | None#

only_diag: bool = False#

out_dim: int | None = None#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

rho: list[Array] | None = None#

rho_t: list[Array] | None = None#

s: Array#

scaled_tau: bool = False#

shape: int#

tau: float | None = None#

timescales(remove_constant=True, n_skip=None)#

property tot_trans#

trans: CvTrans | None = None#

verbose: bool = True#

w: list[Array] | None = None#

w_t: list[Array] | None = None#

weighted_model(chunk_size=None, macro_chunk=1000, out_dim=None, **kwargs)#

Return type:: KoopmanModel

class IMLCV.base.rounds.RoundInformation(round, valid, num, num_vals, tic, folder, name_bias=None, name_md=None)#

Bases: object

folder: Path#

get_bias()#

Return type:: Bias

get_engine()#

Return type:: MDEngine

name_bias: str | None = None#

name_md: str | None = None#

num: int#

num_vals: Array#

round: int#

tic: StaticMdInfo#

valid: bool#

class IMLCV.base.rounds.Rounds(folder)#

Bases: object

Class that bundle all the information in a folder structure. The structure is shown below. Files within parentheses are not always present optional.

folder/

cv_0/

cv.json round 0/

static_trajectory info.h5 bias.json engine.json (invalid) md_0/

trajectory_info.h5 bias.json (bias_new.json) (invalid) (finished)

md_1

…

round 1/: …

…

cv_1/

…

property P#

property T#

add_cv(collective_variable, c=None)#

add_round(bias, stic=None, mde=None, c=None, r=None)#

continue_run(steps, cv_round=None, round=None, plot=True, wait_for_plots=True)#

static create(folder='output', copy=True, new_folder=True)#

Return type:: Rounds

property cv#

data_loader(num=4, out=-1, split_data=False, new_r_cut=-1.0, cv_round=None, ignore_invalid=False, md_trajs=None, start=None, stop=None, time_series=False, T_max_over_T=50, chunk_size=None, get_colvar=True, min_traj_length=None, recalc_cv=False, get_bias_list=False, num_cv_rounds=1, only_finished=True, uniform=True, lag_n=1, colvar=None, check_dtau=True, verbose=False, weight=True, macro_chunk=2000, macro_chunk_nl=5000, only_update_nl=False, n_max=100000.0, wham=True, scale_times=False, reweight_to_fes=False, reweight_inverse_bincount=True, output_FES_bias=False, weighing_method='WHAM', samples_per_bin=10, min_samples_per_bin=3, load_weight=False)#

finish_data(c=None, r=None, i=None)#

folder: Path#

full_path(name)#

Return type:: str

get_bias(c=None, r=None, i=None)#

Return type:: Bias

get_collective_variable(c=None)#

Return type:: CollectiveVariable

get_engine(c=None, r=None)#

Return type:: MDEngine

get_round(c=None)#

get_static_trajectory_info(c=None, r=None)#

invalidate_data(c=None, r=None, i=None)#

is_finished(c=None, r=None, i=None)#

is_valid(c=None, r=None, i=None)#

is_zipped(c, r)#

Return type:: bool | None

iter(start=None, stop=None, num=3, ignore_invalid=False, only_finished=True, c=None, md_trajs=None, print_timings=False)#

Return type:: Iterable[tuple[RoundInformation, TrajectoryInformation]]

iter_ase_atoms(r=None, c=None, num=3, r_cut=None, minkowski_reduce=True, only_finished=False, ignore_invalid=False)#

n(c=None, r=None)#

path(c=None, r=None, i=None)#

Return type:: Path

plot_cv_discovery(start, end, additional_collective_variables=None, additional_collective_variable_names=None, additional_collective_variable_titles=None, plot_biases=True, ignore_invalid=False, only_finished=False, get_fes_bias_kwargs={}, plot_kwargs={})#

static plot_md_run(rnds, i, r, c, fut=None, inputs=[], outputs=[])#

plot_round(c=None, r=None, name_bias=None, name_points=None, dlo_kwargs={}, plot_kwargs={}, plot_points=True, plot_fes=True)#

rel_path(name)#

property round#

run(bias, steps)#

static run_md(steps, sp, inputs=[], outputs=[])#

run_par(biases, steps, plot=True, KEY=42, sp0=None, ignore_invalid=False, md_trajs=None, cv_round=None, wait_for_plots=False, min_traj_length=None, recalc_cv=False, only_finished=True, profile=False, chunk_size=None, macro_chunk=2000, lag_n=20, use_common_bias=True)#

static_trajectory_information(c=None, r=None)#

Return type:: StaticMdInfo

transform_CV(cv_trans, dlo_kwargs=None, dlo=None, chunk_size=None, cv_round_from=None, cv_round_to=None, new_r_cut=None, plot=True, vmax=0.03808799176039228, verbose=True)#

unzip_cv(cv)#

unzip_cv_round(c, r)#

update_CV(transformer, dlo_kwargs=None, dlo=None, chunk_size=None, plot=True, new_r_cut=None, save_samples=True, save_multiple_cvs=False, jac=<function jacrev>, cv_round_from=None, cv_round_to=None, test=False, max_bias=None, transform_bias=True, samples_per_bin=5, min_samples_per_bin=1, percentile=0.1, use_executor=True, n_max=100000.0, vmax=0.03808799176039228, macro_chunk=1000, macro_chunk_nl=5000, verbose=False)#

validate_data(c=None, r=None, i=None)#

write_xyz(c=None, r=None, num=1, repeat=None, minkowski_reduce=True, r_cut=None, only_finished=False, ext='xyz')#

zip_cv(cv)#

zip_cv_round(c, r)#

zip_cv_rounds(begin=1, end=-2)#

class IMLCV.base.rounds.TrajectoryInformation(ti, round, num, folder, name_bias=None, valid=True, finished=False)#

Bases: object

finished: bool = False#

folder: Path#

get_bias()#

Return type:: Bias | None

name_bias: str | None = None#

num: int#

round: int#

ti: TrajectoryInfo#

valid: bool = True#

class IMLCV.base.rounds.WeightOutput(*, weights, p_select, time_scaling=None, bin_counts=None, grid_nums=None, FES_bias=None, FES_bias_std=None, labels=None)#

Bases: MyPyTreeNode

FES_bias: Bias | None = None#

FES_bias_std: Bias | None = None#

bin_counts: list[Array] | None = None#

grid_nums: list[Array] | None = None#

labels: list[int] | None = None#

p_select: list[Array]#

replace(**updates)#: Returns a new object replacing the specified fields with new values.

time_scaling: list[Array] | None = None#

weights: list[Array]#

IMLCV.base package

Contents

IMLCV.base package#

Submodules#

IMLCV.base.CV module#

Args:#

IMLCV.base.CVDiscovery module#

IMLCV.base.MdEngine module#

IMLCV.base.Observable module#

IMLCV.base.UnitsConstants module#

IMLCV.base.bias module#

IMLCV.base.datastructures module#

IMLCV.base.rounds module#

Module contents#