IMLCV.base.CV

Module Contents

Classes

SystemParams
NeighbourList
CV
CvMetric	class to keep track of topology of given CV. Identifies the periodicitie of CVs and maps to unit square with correct peridicities
CvFunInput
CvFunBase
CvFun
CvFunNn	used to instantiate flax linen CvTrans
CvFunDistrax	creates bijective CV function based on a distrax flow. The seup function should initialize the bijector
CombinedCvFun
CvTrans	f can either be a single CV tranformation or a list of transformations
CvTransNN	Base class for all neural network modules. Layers and models should subclass this class.
NormalizingFlow	normalizing flow. _ProtoCvTransNN are stored separately because they need to be initialized by this module in setup
CvFlow
CollectiveVariable

class IMLCV.base.CV.SystemParams

property batched

property shape

coordinates: jax.Array

cell: Array | None

__post_init__()

__getitem__(slices) → SystemParams

__iter__()

__add__(other)

__str__()

Return str(self).

batch() → SystemParams

unbatch() → SystemParams

angles(deg=True) → jax.Array

_get_neighbour_list(r_cut, r_skin: float = 0.0, z_array: tuple[int] | None = None, z_unique: tuple[int] | None = None, num_z_unique: tuple[int] | None = None, num_neighs: int | None = None, nxyz: tuple[int] | None = None) → tuple[bool, NeighbourList | None]

get_neighbour_list(r_cut, z_array: list[int] | Array, r_skin=0.0) → NeighbourList | None

minkowski_reduce() → tuple[SystemParams, jax.Array]

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

wrap_positions(min=False) → tuple[SystemParams, jax.Array]

wrap pos to lie within unit cell

static _wrap_pos(cell: Array | None, coordinates: jax.Array, min=False) → tuple[jax.Array, jax.Array]

canoncialize(min=False) → tuple[SystemParams, jax.Array, jax.Array]

min_distance(index_1, index_2)

class IMLCV.base.CV.NeighbourList

property batched

property shape

property num_neigh

r_cut: jax_dataclasses.Static[jax.numpy.floating]

atom_indices: jax.Array

op_cell: Array | None

op_coor: Array | None

op_center: Array | None

r_skin: jax_dataclasses.Static[jax.numpy.floating]

sp_orig: SystemParams | None

ijk_indices: Array | None

nxyz: jax_dataclasses.Static[tuple[int] | None]

z_array: jax_dataclasses.Static[tuple[int] | None]

z_unique: jax_dataclasses.Static[tuple[int] | None]

num_z_unique: jax_dataclasses.Static[tuple[int] | None]

_pos(sp_orig)

apply_fun_neighbour(sp: SystemParams, func, r_cut=None, fill_value=0, reduce='full', split_z=False, exclude_self=False)

apply_fun_neighbour_pair(sp: SystemParams, func_double, func_single=None, r_cut=None, fill_value=0, reduce='full', split_z=False, exclude_self=True, unique=True)

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,

),

__getitem__(slices)

update(sp: SystemParams) → tuple[bool, NeighbourList]

nl_split_z(p)

static match_kernel(p1: jax.Array, p2: jax.Array, nl1: NeighbourList, nl2: NeighbourList, matching='REMatch', alpha=0.01, mode='divergence', jit=True)

__add__(other)

static stack(*nls: NeighbourList) → NeighbourList

class IMLCV.base.CV.CV

property batched

property batch_dim

property dim

property size

property shape

property combine_dims

property stack_dims

cv: jax.Array

mapped: jax_dataclasses.Static[bool] = False

atomic: jax_dataclasses.Static[bool] = False

_combine_dims: jax_dataclasses.Static[list | None]

_stack_dims: jax_dataclasses.Static[list | None]

__add__(other) → CV

__radd__(other) → CV

__sub__(other) → CV

__rsub__(other) → CV

__mul__(other) → CV

__rmul__(other) → CV

__matmul__(other) → CV

__rmatmul__(other) → CV

__div__(other) → CV

__rdiv__(other) → CV

batch() → CV

__iter__()

__getitem__(idx)

unbatch() → CV

static stack(*cvs: CV) → CV

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

unstack() → list[CV]

split(flatten=False) → list[CV]

inverse operation of combine

static combine(*cvs: CV, flatten=False) → CV

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

class IMLCV.base.CV.CvMetric(periodicities, bounding_box=None)

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

property ndim

norm(x1: CV, x2: CV, k=1.0)

periodic_wrap(x: CV, min=False) → CV

difference(x1: CV, x2: CV) → jax.Array

min_cv(cv: jax.Array)

__periodic_wrap(xs: jax.Array, min=False)

Translate cvs such over unit cell.

min=True calculates distances, False translates one vector inside box

map(x: jax.Array, displace=True) → jax.Array

transform CVs to lie in unit square.

unmap(x: jax.Array, displace=True) → jax.Array

transform CVs to lie in unit square.

__add__(other)

grid(n, endpoints=None, margin=None)

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

class IMLCV.base.CV.CvFunInput

input: int

conditioners: list[int] | None

split(x: CV)

combine(x: CV, res: CV)

class IMLCV.base.CV.CvFunBase

_: dataclasses.KW_ONLY

cv_input: CvFunInput | None

calc(x: CV, nl: NeighbourList | None, reverse=False, log_det=False) → tuple[CV, Array | None]

abstract _calc(x: CV, nl: NeighbourList | None, reverse=False, conditioners: list[CV] | None = None) → CV

_log_Jf(x: CV, nl: NeighbourList | None, reverse=False, conditioners: list[CV] | None = None) → tuple[CV, Array | None]

naive automated implementation, overrride this

class IMLCV.base.CV.CvFun

Bases: CvFunBase

forward: Callable[[CV, NeighbourList | None, CV | None], CV] | None

backward: Callable[[CV, NeighbourList | None, CV | None], CV] | None

_calc(x: CV, nl: NeighbourList | None, reverse=False, conditioners: list[CV] | None = None) → CV

class IMLCV.base.CV.CvFunNn

Bases: flax.linen.Module, CvFunBase

used to instantiate flax linen CvTrans

abstract setup()

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:

1. Immediately when invoking apply(), init() or init_and_output().

2. 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.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

_calc(x: CV, nl: NeighbourList | None, reverse=False, conditioners: list[CV] | None = None) → CV

abstract forward(x: CV, nl: NeighbourList | None, conditioners: list[CV] | None = None) → CV

abstract backward(x: CV, nl: NeighbourList | None, conditioners: list[CV] | None = None) → CV

class IMLCV.base.CV.CvFunDistrax

Bases: flax.linen.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

abstract setup()

setups self.bijector

_calc(x: CV, nl: NeighbourList | None, reverse=False, log_det=False, conditioners: list[CV] | None = None) → CV

_log_Jf(x: CV, nl: NeighbourList | None, reverse=False, conditioners: list[CV] | None = None) → tuple[CV, Array | None]

naive implementation, overrride this

class IMLCV.base.CV.CombinedCvFun

Bases: CvFunBase

classes: list[list[CvFunBase]]

calc(x: CV, nl: NeighbourList | None, reverse=False, log_det=False) → tuple[CV, Array | None]

abstract _log_Jf(x: CV, nl: NeighbourList | None, reverse=False, conditioners: list[CV] | None = None) → tuple[CV, Array | None]

naive automated implementation, overrride this

class IMLCV.base.CV.CvTrans

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

trans: list[CvFunBase]

static from_array_function(f: collections.abc.Callable[[jax.Array, NeighbourList | None, None], jax.Array])

static from_cv_function(f: collections.abc.Callable[[CV, NeighbourList | None, CV | None], CV]) → CvTrans

static from_cv_fun(proto: CvFunBase)

compute_cv_trans(x: CV, nl: NeighbourList | None = None, reverse=False, log_Jf=False) → tuple[CV, Array | None]

result is always batched arg: CV

__mul__(other)

__add__(other: CvTrans) → CvTrans

static stack(*cv_trans: CvTrans)

class IMLCV.base.CV.CvTransNN

Bases: flax.linen.Module, CvTrans

Base class for all neural network modules. Layers and models should subclass this class.

All Flax Modules are Python 3.7 dataclasses. Since dataclasses take over __init__, you should instead override setup(), which is automatically called to initialize the module.

Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the setup() method.

You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased, __call__ is a popular choice because it allows you to use module instances as if they are functions:

from flax import linen as nn

class Module(nn.Module):
  features: Tuple[int, ...] = (16, 4)

  def setup(self):
    self.dense1 = Dense(self.features[0])
    self.dense2 = Dense(self.features[1])

  def __call__(self, x):
    return self.dense2(nn.relu(self.dense1(x)))

Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the compact() wrapper.

trans: list[CvFunBase]

setup() → None

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:

1. Immediately when invoking apply(), init() or init_and_output().

2. 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.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

compute_cv_trans(x: CV, nl: NeighbourList | None, reverse=False, log_Jf=False) → tuple[CV, Array | None]

result is always batched arg: CV

__mul__(other)

class IMLCV.base.CV.NormalizingFlow

Bases: flax.linen.Module

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

flow: CvTransNN | CvTransNN

setup() → None

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:

1. Immediately when invoking apply(), init() or init_and_output().

2. 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.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

calc(x: CV, nl: NeighbourList | None, reverse: bool, test_log_det=False)

class IMLCV.base.CV.CvFlow(func: collections.abc.Callable[[SystemParams, NeighbourList | None], CV], trans: CvTrans | None = None)

static from_function(f: collections.abc.Callable[[SystemParams, NeighbourList | None], jax.Array], atomic=False) → CvFlow

compute_cv_flow(x: SystemParams, nl: NeighbourList | None = None, jit=True, chunk_size: int | None = None) → CV

__add__(other)

__mul__(other)

find_sp(x0: SystemParams, target: CV, target_nl: NeighbourList, nl0: NeighbourList | None = None, maxiter=10000, tol=0.0001, norm=lambda cv1, cv2, nl1, nl2: ..., solver=jaxopt.GradientDescent) → SystemParams

class IMLCV.base.CV.CollectiveVariable(f: CvFlow, metric: CvMetric, jac=jacrev)

property n

_jit_f(sp, nl)

_jit_df(sp, nl)

compute_cv(sp: SystemParams, nl: NeighbourList | None = None, jacobian=False, jit=True, chunk_size: int | None = None) → tuple[CV, CV]