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)[source]#

Bases: object

atomic: bool = False[source]#

batch()[source]#

Return type:: CV

property batch_dim[source]#

property batched[source]#

static combine(*cvs, flatten=False)[source]#

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[source]#

cv: Array[source]#

property dim[source]#

mapped: bool = False[source]#

property shape[source]#

property size[source]#

split(flatten=False)[source]#

inverse operation of combine

Return type:: list[CV]

static stack(*cvs)[source]#

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[source]#

unbatch()[source]#

Return type:: CV

unstack()[source]#

Return type:: list[CV]

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

Bases: object

compute_cv(sp, nl=None, jacobian=False, jit=True, chunk_size=None)[source]#

Return type:: tuple[CV, CV]

property n[source]#

class IMLCV.base.CV.CombinedCvFun(*, cv_input=None, classes)[source]#

Bases: CvFunBase

calc(x, nl, reverse=False, log_det=False)[source]#

Return type:: tuple[CV, Optional[Array]]

classes: list[list[CvFunBase]][source]#

class IMLCV.base.CV.CvFlow(func, trans=None)[source]#

Bases: object

compute_cv_flow(x, nl=None, jit=True, chunk_size=None)[source]#

Return type:: CV

find_sp(x0, target, target_nl, nl0=None, maxiter=10000, tol=0.0001, norm=<function CvFlow.<lambda>>, solver=<class 'jaxopt._src.gradient_descent.GradientDescent'>)[source]#

Return type:: SystemParams

static from_function(f, atomic=False)[source]#

Return type:: CvFlow

class IMLCV.base.CV.CvFun(*, cv_input=None, forward=None, backward=None)[source]#

Bases: CvFunBase

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

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

class IMLCV.base.CV.CvFunBase(*, cv_input=None)[source]#

Bases: object

calc(x, nl, reverse=False, log_det=False)[source]#

Return type:: tuple[CV, Optional[Array]]

cv_input: CvFunInput | None = None[source]#

class IMLCV.base.CV.CvFunDistrax(parent=<flax.linen.module._Sentinel object>, name=None, *, cv_input=None)[source]#

Bases: Module, CvFunBase

creates bijective CV function based on a distrax flow. The seup function should initialize the bijector

class RealNVP(CvFunDistrax):

_: dataclasses.KW_ONLY latent_dim: int

def setup(self):

self.s = Dense(features=self.latent_dim) self.t = Dense(features=self.latent_dim)

# Alternating binary mask. self.bijector = distrax.as_bijector(

tfp.bijectors.RealNVP(
fraction_masked=0.5, shift_and_log_scale_fn=self.shift_and_scale,

)

)

def shift_and_scale(self, x0, input_depth, **condition_kwargs):

return self.s(x0), self.t(x0)

bijector: distrax.Bijector[source]#

name: Optional[str] = None#

parent: Union[Type[Module], Type[Scope], Type[_Sentinel], None] = None#

scope: Optional[Scope] = None#

abstract setup()[source]#: setups self.bijector

class IMLCV.base.CV.CvFunInput(*, input, conditioners=None)[source]#

Bases: object

combine(x, res)[source]#

conditioners: Optional[list[int]] = None[source]#

input: int[source]#

split(x)[source]#

class IMLCV.base.CV.CvFunNn(parent=<flax.linen.module._Sentinel object>, name=None, *, cv_input=None)[source]#

Bases: Module, CvFunBase

used to instantiate flax linen CvTrans

abstract backward(x, nl, conditioners=None)[source]#

Return type:: CV

abstract forward(x, nl, conditioners=None)[source]#

Return type:: CV

name: Optional[str] = None#

parent: Union[Type[Module], Type[Scope], Type[_Sentinel], None] = None#

scope: Optional[Scope] = None#

abstract setup()[source]#

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

Immediately when invoking apply(), init() or init_and_output().
Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):
class MyModule(nn.Module):
  def setup(self):
    submodule = Conv(...)

    # Accessing `submodule` attributes does not yet work here.

    # The following line invokes `self.__setattr__`, which gives
    # `submodule` the name "conv1".
    self.conv1 = submodule

    # Accessing `submodule` attributes or methods is now safe and
    # either causes setup() to be called once.
Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class IMLCV.base.CV.CvMetric(periodicities, bounding_box=None)[source]#

Bases: object

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

difference(x1, x2)[source]#

Return type:: Array

grid(n, endpoints=None, margin=None)[source]#

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(x, displace=True)[source]#

transform CVs to lie in unit square.

Return type:: Array

min_cv(cv)[source]#

property ndim[source]#

norm(x1, x2, k=1.0)[source]#

periodic_wrap(x, min=False)[source]#

Return type:: CV

unmap(x, displace=True)[source]#

transform CVs to lie in unit square.

Return type:: Array

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

Bases: object

f can either be a single CV tranformation or a list of transformations

compute_cv_trans(x, nl=None, reverse=False, log_Jf=False)[source]#

result is always batched arg: CV

Return type:: tuple[CV, Optional[Array]]

static from_array_function(f)[source]#

static from_cv_fun(proto)[source]#

static from_cv_function(f)[source]#

Return type:: CvTrans

static stack(*cv_trans)[source]#

trans: list[CvFunBase][source]#

class IMLCV.base.CV.CvTransNN(trans, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Bases: Module, CvTrans

compute_cv_trans(x, nl, reverse=False, log_Jf=False)[source]#

result is always batched arg: CV

Return type:: tuple[CV, Optional[Array]]

name: Optional[str] = None#

parent: Union[Type[Module], Type[Scope], Type[_Sentinel], None] = None#

scope: Optional[Scope] = None#

setup()[source]#

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases: :rtype: None

Immediately when invoking apply(), init() or init_and_output().
Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):
class MyModule(nn.Module):
  def setup(self):
    submodule = Conv(...)

    # Accessing `submodule` attributes does not yet work here.

    # The following line invokes `self.__setattr__`, which gives
    # `submodule` the name "conv1".
    self.conv1 = submodule

    # Accessing `submodule` attributes or methods is now safe and
    # either causes setup() to be called once.
Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

trans: list[CvFunBase][source]#

class IMLCV.base.CV.NeighbourList(r_cut, atom_indices, op_cell, op_coor, op_center, r_skin, sp_orig=None, ijk_indices=None, nxyz=None, z_array=None, z_unique=None, num_z_unique=None)[source]#

Bases: object

apply_fun_neighbour(sp, func, r_cut=None, fill_value=0, reduce='full', split_z=False, exclude_self=False)[source]#

apply_fun_neighbour_pair(sp, func_double, func_single=None, r_cut=None, fill_value=0, reduce='full', split_z=False, exclude_self=True, unique=True)[source]#: Args:#

func_single=lambda r_ij, atom_index_j: (1,), func_double=lambda r_ij, atom_index_j, data_j, r_ik, atom_index_k, data_k: (

r_ij, atom_index_j, r_ik, atom_index_k,

),

atom_indices: Array[source]#

property batched[source]#

ijk_indices: Optional[Array] = None[source]#

static match_kernel(p1, p2, nl1, nl2, matching='REMatch', alpha=0.01, mode='divergence', jit=True)[source]#

nl_split_z(p)[source]#

property num_neigh[source]#

num_z_unique: Optional[tuple[int]] = None[source]#

nxyz: Optional[tuple[int]] = None[source]#

op_cell: Optional[Array][source]#

op_center: Optional[Array][source]#

op_coor: Optional[Array][source]#

r_cut: floating[source]#

r_skin: floating[source]#

property shape[source]#

sp_orig: Optional[SystemParams] = None[source]#

static stack(*nls)[source]#

Return type:: NeighbourList

update(sp)[source]#

Return type:: tuple[bool, NeighbourList]

z_array: Optional[tuple[int]] = None[source]#

z_unique: Optional[tuple[int]] = None[source]#

class IMLCV.base.CV.NormalizingFlow(flow, parent=<flax.linen.module._Sentinel object>, name=None)[source]#

Bases: Module

normalizing flow. _ProtoCvTransNN are stored separately because they need to be initialized by this module in setup

calc(x, nl, reverse, test_log_det=False)[source]#

flow: CvTransNN | CvTransNN[source]#

name: Optional[str] = None#

parent: Union[Type[Module], Type[Scope], Type[_Sentinel], None] = None#

scope: Optional[Scope] = None#

setup()[source]#

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases: :rtype: None

Immediately when invoking apply(), init() or init_and_output().
Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):
class MyModule(nn.Module):
  def setup(self):
    submodule = Conv(...)

    # Accessing `submodule` attributes does not yet work here.

    # The following line invokes `self.__setattr__`, which gives
    # `submodule` the name "conv1".
    self.conv1 = submodule

    # Accessing `submodule` attributes or methods is now safe and
    # either causes setup() to be called once.
Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class IMLCV.base.CV.SystemParams(coordinates, cell=None)[source]#

Bases: object

angles(deg=True)[source]#

Return type:: Array

batch()[source]#

Return type:: SystemParams

property batched[source]#

canoncialize(min=False)[source]#

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

cell: Optional[Array] = None[source]#

coordinates: Array[source]#

get_neighbour_list(r_cut, z_array, r_skin=0.0)[source]#

Return type:: Optional[NeighbourList]

min_distance(index_1, index_2)[source]#

minkowski_reduce()[source]#

base on code from ASE: https://wiki.fysik.dtu.dk/ase/_modules/ase/geometry/minkowski_reductsp.cellion.html#minkowski_reduce

Return type:: tuple[SystemParams, Array]

property shape[source]#

unbatch()[source]#

Return type:: SystemParams

wrap_positions(min=False)[source]#

wrap pos to lie within unit cell

Return type:: tuple[SystemParams, Array]

IMLCV.base.CVDiscovery module#

class IMLCV.base.CVDiscovery.CVDiscovery(transformer)[source]#

Bases: object

convert set of coordinates to good collective variables.

compute(rounds, num_rounds=4, samples=3000.0, plot=True, new_r_cut=None, chunk_size=None, split_data=False, name=None, **kwargs)[source]#

Return type:: CollectiveVariable

data_loader(rounds, num=4, out=-1, split_data=False, new_r_cut=None)[source]#

Return type:: tuple[list[SystemParams], Optional[list[NeighbourList]], CollectiveVariable, StaticTrajectoryInfo]

static plot_app(sps, nl, old_cv, new_cv, name, labels=None, chunk_size=None)[source]#

class IMLCV.base.CVDiscovery.Transformer(outdim, periodicity=None, bounding_box=None, descriptor='sb', descriptor_kwargs={})[source]#

Bases: object

fit(sp_list, nl_list, chunk_size=None, prescale=True, postscale=True, jac=<function jacrev>, *fit_args, **fit_kwargs)[source]#

Return type:: tuple[CV, CollectiveVariable]

post_fit(y, scale)[source]#

Return type:: tuple[CV, CvTrans]

pre_fit(z, nl, chunk_size=None, scale=True)[source]#

Return type:: tuple[list[CV], CvFlow]

IMLCV.base.MdEngine module#

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

Currently, the MD is done with YAFF/OpenMM

class IMLCV.base.MdEngine.MDEngine(bias, energy, static_trajectory_info, trajectory_file=None, sp=None)[source]#

Bases: ABC

Base class for MD engine.

class YaffSys(_ener, _tic)[source]#

Bases: object

class YaffCell(_ener)[source]#

Bases: object

property nvec[source]#

property rvecs[source]#

update_rvecs(rvecs)[source]#

property volume[source]#

property cell[source]#

property charges[source]#

property masses[source]#

property natom[source]#

property numbers[source]#

property pos[source]#

get_bias(gpos=False, vtens=False)[source]#

Return type:: tuple[CV, EnergyResult]

get_energy(gpos=False, vtens=False)[source]#

Return type:: EnergyResult

get_trajectory()[source]#

Return type:: TrajectoryInfo

keys = ['bias', 'energy', 'static_trajectory_info', 'trajectory_file'][source]#

static load(file, **kwargs)[source]#

Return type:: MDEngine

new_bias(bias, **kwargs)[source]#

Return type:: MDEngine

property nl: NeighbourList | None[source]#

run(steps)[source]#

run the integrator for a given number of steps.

Parameters:: steps – number of MD steps

save(file)[source]#

save_step(T=None, P=None, t=None, err=None)[source]#

property sp: SystemParams[source]#

property yaff_system: YaffSys[source]#

class IMLCV.base.MdEngine.StaticTrajectoryInfo(timestep, T, timecon_thermo, atomic_numbers, r_cut=None, P=None, timecon_baro=None, write_step=100, equilibration=None, screen_log=1000, max_grad=0.04031059429903557)[source]#

Bases: object

P: Optional[float] = None[source]#

T: float[source]#

atomic_numbers: Array[source]#

property barostat[source]#

equilibration: Optional[float] = None[source]#

static load(filename)[source]#

Return type:: StaticTrajectoryInfo

property masses[source]#

max_grad: Optional[float] = 0.04031059429903557[source]#

r_cut: Optional[float] = None[source]#

save(filename)[source]#

screen_log: int = 1000[source]#

property thermostat[source]#

timecon_baro: Optional[float] = None[source]#

timecon_thermo: float[source]#

timestep: float[source]#

write_step: int = 100[source]#

class IMLCV.base.MdEngine.TrajectoryInfo(_positions, _cell=None, _charges=None, _e_pot=None, _e_pot_gpos=None, _e_pot_vtens=None, _e_bias=None, _e_bias_gpos=None, _e_bias_vtens=None, _cv=None, _T=None, _P=None, _err=None, _t=None, _capacity=-1, _size=-1)[source]#

Bases: object

property CV: CV | None[source]#

property P: Array | None[source]#

property T: Array | None[source]#

property cell: Array | None[source]#

property charges: Array | None[source]#

property cv: Array | None[source]#

property e_bias: Array | None[source]#

property e_bias_gpos: Array | None[source]#

property e_bias_vtens: Array | None[source]#

property e_pot: Array | None[source]#

property e_pot_gpos: Array | None[source]#

property e_pot_vtens: Array | None[source]#

property err: Array | None[source]#

static load(filename)[source]#

Return type:: TrajectoryInfo

property positions: Array | None[source]#

save(filename)[source]#

property shape[source]#

property sp: SystemParams[source]#

property t: Array | None[source]#

property volume[source]#

IMLCV.base.Observable module#

class IMLCV.base.Observable.ThermoLIB(rounds, rnd=None)[source]#

Bases: object

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

fes_bias(plot=True, max_bias=None, fs=None, choice='gridbias', n=None, start_r=0, rbf_kernel='thin_plate_spline', rbf_degree=None, smoothing_threshold=0.0019043995880196139, samples_per_bin=500, **plot_kwargs)[source]#

fes_nd_thermolib(plot=True, n=None, start_r=0, update_bounding_box=False, samples_per_bin=500)[source]#

new_metric(plot=False, r=None)[source]#

time_per_bin = 82682.74667329364[source]#

IMLCV.base.bias module#

class IMLCV.base.bias.BC[source]#

Bases: object

base class for biased Energy of MD simulation.

static load(filename)[source]#

Return type:: BC

save(filename)[source]#

class IMLCV.base.bias.Bias(collective_variable, start=None, step=None)[source]#

Bases: BC, ABC

base class for biased MD runs.

compute_from_cv(cvs, diff=False, jit=True)[source]#

compute the energy and derivative.

If map==False, the cvs are assumed to be already mapped

Return type:: CV

compute_from_system_params(sp, gpos=False, vir=False, nl=None, jit=True)[source]#

Computes the bias, the gradient of the bias wrt the coordinates and the virial.

Return type:: tuple[CV, EnergyResult]

finalize()[source]#

Should be called at end of metadynamics simulation.

Optimises compute

abstract get_args()[source]#: function that return dictionary with kwargs of _compute.

static load(filename)[source]#

Return type:: Bias

plot(name, x_unit=None, y_unit=None, n=50, traj=None, vmin=0, vmax=0.03808799176039228, map=False, inverted=False, margin=None, x_lim=None, y_lim=None, bins=None)[source]#: plot bias.