IMLCV.base.rounds

Attributes

X
X2

Classes

TrajectoryInformation
RoundInformation
Rounds	Class that bundle all the information in a folder structure. The structure is shown below. Files within parentheses are not always present optional.
WeightOutput	Base class for dataclasses that should act like a JAX pytree node.
DataLoaderOutput	Base class for dataclasses that should act like a JAX pytree node.
KoopmanModel	Base class for dataclasses that should act like a JAX pytree node.
Covariances	Base class for dataclasses that should act like a JAX pytree node.

Module Contents

class IMLCV.base.rounds.TrajectoryInformation

ti: IMLCV.base.MdEngine.TrajectoryInfo

round: int

num: int

folder: pathlib.Path

name_bias: str | None = None

valid: bool = True

finished: bool = False

get_bias() → IMLCV.base.bias.Bias | None

class IMLCV.base.rounds.RoundInformation

round: int

valid: bool

num: int

num_vals: jax.Array

tic: IMLCV.base.MdEngine.StaticMdInfo

folder: pathlib.Path

name_bias: str | None = None

name_md: str | None = None

get_bias() → IMLCV.base.bias.Bias

get_engine() → IMLCV.base.MdEngine.MDEngine

class IMLCV.base.rounds.Rounds

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/

…

…

folder: pathlib.Path

static create(folder: str | pathlib.Path = 'output', copy=True, new_folder=True) → Rounds

__getstate__()

__setstate__(state)

full_path(name: str | pathlib.Path) → str

rel_path(name: str | pathlib.Path)

path(c: int | None = None, r: int | None = None, i: int | None = None) → pathlib.Path

is_zipped(c, r) → bool | None

zip_cv_round(c, r)

unzip_cv_round(c, r)

zip_cv(cv)

unzip_cv(cv)

zip_cv_rounds(begin=1, end=-2)

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

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

plot_cv_discovery(start: int, end: int, additional_collective_variables: list[IMLCV.base.CV.CollectiveVariable] | None = None, additional_collective_variable_names: list[str] | None = None, additional_collective_variable_titles: list[str] | None = None, plot_biases=True, ignore_invalid=False, only_finished=False, get_fes_bias_kwargs={}, plot_kwargs={})

_c_vals() → list[int]

_r_vals(c: int | None = None)

_i_vals(c: int | None = None, r: int | None = None)

_name_md(c, r)

_name_bias(c, r, i: int | None = None)

_num_vals(c, r: int | None = None)

add_cv(collective_variable: IMLCV.base.CV.CollectiveVariable, c: int | None = None)

add_round(bias: IMLCV.base.bias.Bias, stic: IMLCV.base.MdEngine.StaticMdInfo | None = None, mde=None, c: int | None = None, r: int | None = None)

iter(start=None, stop=None, num=3, ignore_invalid=False, only_finished=True, c: int | None = None, md_trajs: list[int] | None = None, print_timings=False) → collections.abc.Iterable[tuple[RoundInformation, TrajectoryInformation]]

data_loader(num: int = 4, out: int = -1, split_data: bool = False, new_r_cut: float | None = -1.0, cv_round: int | None = None, ignore_invalid=False, md_trajs: list[int] | None = None, start: int | None = None, stop: int | None = None, time_series: bool = False, T_max_over_T=50, chunk_size: int | None = None, get_colvar: bool = True, min_traj_length: int | None = None, recalc_cv: bool = False, get_bias_list: bool = False, num_cv_rounds: int = 1, only_finished: bool = True, uniform: bool = True, lag_n: int = 1, colvar: IMLCV.base.CV.CollectiveVariable | None = None, check_dtau=True, verbose: bool = False, weight: bool = True, macro_chunk: int | None = 2000, macro_chunk_nl: int | None = 5000, only_update_nl: bool = False, n_max: int | float = 100000.0, wham: bool = True, scale_times: bool = False, reweight_to_fes: bool = False, reweight_inverse_bincount: bool = True, output_FES_bias: bool = False, weighing_method: str = 'WHAM', samples_per_bin: int = 10, min_samples_per_bin: int = 3, load_weight: bool = False)

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

_trajectory_information(r: int, i: int, c: int | None = None) → TrajectoryInformation

static_trajectory_information(c: int | None = None, r: int | None = None) → IMLCV.base.MdEngine.StaticMdInfo

_round_information(c: int | None = None, r: int | None = None) → RoundInformation

property T

property P

property round

get_round(c: int | None = None)

property cv

n(c: int | None = None, r: int | None = None)

invalidate_data(c: int | None = None, r: int | None = None, i: int | None = None)

validate_data(c: int | None = None, r: int | None = None, i: int | None = None)

finish_data(c: int | None = None, r: int | None = None, i: int | None = None)

is_valid(c: int | None = None, r: int | None = None, i: int | None = None)

is_finished(c: int | None = None, r: int | None = None, i: int | None = None)

get_collective_variable(c: int | None = None) → IMLCV.base.CV.CollectiveVariable

get_static_trajectory_info(c: int | None = None, r: int | None = None)

get_bias(c: int | None = None, r: int | None = None, i: int | None = None) → IMLCV.base.bias.Bias

get_engine(c: int | None = None, r: int | None = None) → IMLCV.base.MdEngine.MDEngine

run(bias, steps)

run_par(biases: Sequence[IMLCV.base.bias.Bias], steps, plot=True, KEY=42, sp0: IMLCV.base.CV.SystemParams | None = None, ignore_invalid=False, md_trajs: list[int] | None = None, cv_round: int | None = 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)

continue_run(steps: int, cv_round: int | None = None, round: int | None = None, plot=True, wait_for_plots=True)

static _get_init(rounds: Rounds, KEY: jax.Array | int, common_bias_name: str | None, biases: Sequence[IMLCV.base.bias.Bias], ignore_invalid: bool = False, only_finished: bool = True, min_traj_length: int | None = None, recalc_cv: bool = False, chunk_size: int | None = None, md_trajs: list[int] | None = None, cv_round: int | None = None, sp0: IMLCV.base.CV.SystemParams | None = None, r: int | None = None, macro_chunk: int | None = 1000, lag_n: int = 20, out: int = 20000, use_energies: bool = False)

static run_md(steps: int, sp: IMLCV.base.CV.SystemParams | None, inputs: list[pathlib.Path] = [], outputs: list[pathlib.Path] = [])

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

update_CV(transformer: IMLCV.base.CVDiscovery.Transformer, dlo_kwargs=None, dlo: DataLoaderOutput | None = None, chunk_size=None, plot=True, new_r_cut=None, save_samples=True, save_multiple_cvs=False, jac=jax.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=100 * kjmol, macro_chunk=1000, macro_chunk_nl: int = 5000, verbose=False)

static _update_CV(rounds: Rounds, transformer: IMLCV.base.CVDiscovery.Transformer, cv_round_from: int, cv_round_to: int, dlo_kwargs={}, dlo: DataLoaderOutput | None = None, chunk_size=None, macro_chunk=1000, macro_chunk_nl: int = 5000, plot=True, new_r_cut=None, save_samples=True, save_multiple_cvs=False, jac=jax.jacrev, test=False, max_bias=None, transform_bias=True, samples_per_bin=5, min_samples_per_bin=1, percentile=0.1, n_max=100000.0, vmax=100 * kjmol, verbose=True)

transform_CV(cv_trans: IMLCV.base.CV.CvTrans, dlo_kwargs=None, dlo: DataLoaderOutput | None = None, chunk_size=None, cv_round_from=None, cv_round_to=None, new_r_cut=None, plot=True, vmax=100 * kjmol, verbose=True)

__update_CV(new_collective_variable: IMLCV.base.CV.CollectiveVariable, new_bias: IMLCV.base.bias.Bias, cv_round_from: int, cvs_new: list[IMLCV.base.CV.CV], dlo: DataLoaderOutput, new_r_cut=None, save_samples=True, save_multiple_cvs=False)

_copy_from_previous_round(dlo: DataLoaderOutput, new_cvs: list[IMLCV.base.CV.CV], invalidate: bool = False, cv_round: int | None = None)

class IMLCV.base.rounds.WeightOutput(*args, **kwargs)

Bases: IMLCV.base.datastructures.MyPyTreeNode

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

weights: list[jax.Array]

p_select: list[jax.Array]

time_scaling: list[jax.Array] | None = None

bin_counts: list[jax.Array] | None = None

grid_nums: list[jax.Array] | None = None

FES_bias: IMLCV.base.bias.Bias | None = None

FES_bias_std: IMLCV.base.bias.Bias | None = None

labels: list[int] | None = None

class IMLCV.base.rounds.DataLoaderOutput(*args, **kwargs)

Bases: IMLCV.base.datastructures.MyPyTreeNode

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

sp: list[IMLCV.base.CV.SystemParams]

cv: list[IMLCV.base.CV.CV]

sti: IMLCV.base.MdEngine.StaticMdInfo

ti: list[IMLCV.base.MdEngine.TrajectoryInfo]

collective_variable: IMLCV.base.CV.CollectiveVariable

labels: list[jax.Array] | None = None

nl: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None

sp_t: list[IMLCV.base.CV.SystemParams] | None = None

nl_t: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None

cv_t: list[IMLCV.base.CV.CV] | None = None

ti_t: list[IMLCV.base.MdEngine.TrajectoryInfo] | None = None

time_series: bool = False

tau: float | None = None

bias: list[IMLCV.base.bias.Bias] | None = None

ground_bias: IMLCV.base.bias.Bias | None = None

_weights: list[jax.Array] | None = None

_weights_t: list[jax.Array] | None = None

_rho: list[jax.Array] | None = None

_rho_t: list[jax.Array] | None = None

scaled_tau: bool = False

__add__(other)

static get_histo(data_nums: list[IMLCV.base.CV.CV], weights: None | list[jax.Array] = None, log_w=False, macro_chunk=320, verbose=False, shape_mask=None, nn=-1, f_func: Callable | None = None) → jax.Array

static _histogram(metric: IMLCV.base.CV.CvMetric, n_grid=40, grid_bounds=None, chunk_size=None, chunk_size_mid=1)

static _unstack_weights(stack_dims, weights: jax.Array) → list[jax.Array]

static norm_w(w_stacked: jax.Array, ret_n=False) → jax.Array | tuple[jax.Array, jax.Array]

static check_w(w_stacked: jax.Array) → jax.Array

koopman_weight(w: list[jax.Array] | None = None, w_t: list[jax.Array] | None = None, samples_per_bin: int = 50, max_bins: int | float = 100000.0, out_dim: int = -1, chunk_size: int | None = None, indicator_CV: bool = True, koopman_eps: float = 1e-05, koopman_eps_pre: float = 0, cv_0: list[IMLCV.base.CV.CV] | None = None, cv_t: list[IMLCV.base.CV.CV] | None = None, macro_chunk: int = 1000, verbose: bool = False, max_features_koopman: int = 5000, margin: float = 0.1, add_1: bool | None = None, only_diag: bool = False, calc_pi: bool = False, sparse: bool = False, output_w_corr: bool = False, correlation: bool = True, return_km: bool = False, labels: jax.Array | list[int] | None = None, koopman_kwargs: dict = {})

wham_weight(samples_per_bin: int = 10, n_max: float | int = 100000.0, chunk_size: int | None = None, wham_eps=1e-07, cv_0: list[IMLCV.base.CV.CV] | None = None, cv_t: list[IMLCV.base.CV.CV] | None = None, macro_chunk: int | None = 1000, verbose: bool = False, margin: float = 0.1, bias_cutoff: float = 300 * kjmol, min_f_i: float = 1e-30, log_sum_exp: bool = True, return_bias: bool = False, return_std_bias: bool = False, min_samples: int = 3, lagrangian: bool = False, max_sigma: float = 5 * kjmol, sparse_inverse: bool = True, inverse_sigma_weighting: bool = False, output_bincount: bool = True, output_free_energy: bool = True, output_time_scaling: bool = True, smooth_bias: bool = False, max_bias_margin: float = 0.1, alpha_dirichlet: float = 1.0) → WeightOutput

get_bincount(cv_0: list[IMLCV.base.CV.CV] | None = None, w: list[jax.Array] | None = 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)

static _transform(cv, nl, shmap, shmap_kwargs, argmask: jax.Array | None = None, pi: jax.Array | None = None, add_1: bool = False, add_1_pre: bool = False, q: jax.Array | None = None, l: jax.Array | None = None) → jax.Array

koopman_model(cv_0: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams] | None = None, cv_t: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams] | None = None, nl: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, nl_t: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, only_return_weights=False, symmetric=False, rho: list[jax.Array] | None = None, w: list[jax.Array] | None = None, rho_t: list[jax.Array] | None = None, w_t: list[jax.Array] | None = 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: float = 1.0, only_diag=False, calc_pi=False, scaled_tau=None, sparse=True, correlation=True, auto_cov_threshold: float | None = None) → KoopmanModel

filter_nans(x: list[IMLCV.base.CV.CV] | None = None, x_t: list[IMLCV.base.CV.CV] | None = None, macro_chunk=1000)

static apply_cv(f: IMLCV.base.CV.CvTrans, x: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams], x_t: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams] | None = None, nl: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, nl_t: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, chunk_size: int | None = None, macro_chunk: int | None = 1000, shmap: bool = True, shmap_kwargs=ShmapKwargs.create(), verbose: bool = False, print_every: int = 10, jit_f: bool = True) → tuple[list[IMLCV.base.CV.CV], list[IMLCV.base.CV.CV] | None]

static _apply(x: list[X], f: Callable[[X, IMLCV.base.CV.NeighbourList | None], X2], x_t: list[X] | None = None, nl: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, nl_t: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, macro_chunk: int | None = 1000, verbose: bool = False, jit_f: bool = True, print_every: int = 10) → tuple[list[X2], list[X2] | None]

static _apply_bias(x: list[IMLCV.base.CV.CV], bias: IMLCV.base.bias.Bias, chunk_size=None, macro_chunk=1000, verbose=False, shmap=False, shmap_kwargs=ShmapKwargs.create(), jit_f=True) → list[jax.Array]

apply_bias(bias: IMLCV.base.bias.Bias, chunk_size=None, macro_chunk=1000, verbose=False, shmap=False) → list[jax.Array]

get_point(point: IMLCV.base.CV.CV, sigma=0.01, key=jax.random.PRNGKey(0)) → tuple[IMLCV.base.CV.CV, IMLCV.base.CV.SystemParams, IMLCV.base.CV.NeighbourList | IMLCV.base.CV.NeighbourListInfo | None, IMLCV.base.MdEngine.TrajectoryInfo]

get_fes_bias_from_weights(cv: list[IMLCV.base.CV.CV] | None = None, weights: list[jax.Array] | None = None, rho: list[jax.Array] | None = None, samples_per_bin=100, min_samples_per_bin: int | None = 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: IMLCV.base.CV.CollectiveVariable | None = None, set_outer_border=True, rbf_degree: int | None = None)

static _get_fes_bias_from_weights(T, weights: list[jax.Array], rho: list[jax.Array], collective_variable: IMLCV.base.CV.CollectiveVariable, cv: list[IMLCV.base.CV.CV], samples_per_bin=100, min_samples_per_bin: int | None = 5, n_max=1000.0, n_grid=None, max_bias=None, chunk_size=None, macro_chunk=1000, max_bias_margin=0.2, rbf_bias=True, kernel='gaussian', set_outer_border=True, rbf_degree: int | None = None) → IMLCV.implementations.bias.RbfBias

get_transformed_fes(new_cv: list[IMLCV.base.CV.CV], new_colvar: IMLCV.base.CV.CollectiveVariable, samples_per_bin=5, min_samples_per_bin: int = 1, chunk_size=1, smoothing=0.0, max_bias=None, shmap=False, n_grid_old=50, n_grid_new=30) → IMLCV.implementations.bias.RbfBias

abstract transform_FES(trans: IMLCV.base.CV.CvTrans, T: float | None = None, max_bias=100 * kjmol, n_grid=25)

recalc(chunk_size: int | None = None, macro_chunk: int | None = 1000, shmap: bool = False, verbose: bool = False)

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

class IMLCV.base.rounds.KoopmanModel(*args, **kwargs)

Bases: IMLCV.base.datastructures.MyPyTreeNode

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

s: jax.Array

cov: Covariances

W0: jax.Array

W1: jax.Array

argmask: jax.Array | None

shape: int

cv_0: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams]

cv_t: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams]

nl: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None

nl_t: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None

w: list[jax.Array] | None = None

rho: list[jax.Array] | None = None

w_t: list[jax.Array] | None = None

rho_t: list[jax.Array] | None = None

eps: float = 1e-10

eps_pre: float | None = None

only_diag: bool = False

calc_pi: bool = True

scaled_tau: bool = False

add_1: bool = False

max_features: int = 5000

max_features_pre: int = 5000

out_dim: int | None = None

correlation_whiten: bool = True

verbose: bool = True

tau: float | None = None

T_scale: float = 1.0

trans: IMLCV.base.CV.CvTrans | IMLCV.base.CV.CvTrans | None = None

constant_threshold: float = 1e-10

static create(w: list[jax.Array] | None, rho: list[jax.Array] | None, w_t: list[jax.Array] | None, rho_t: list[jax.Array] | None, cv_0: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams], cv_t: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams], nl: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, nl_t: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, add_1=True, eps: float = 1e-14, eps_pre: float = 0, symmetric=False, out_dim=-1, max_features=5000, max_features_pre=5000, tau=None, macro_chunk=1000, chunk_size=None, verbose=True, trans: IMLCV.base.CV.CvTrans | None = None, T_scale: float = 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: float = 1e-10)

static _add_1(cv, nl, shmap, shmap_kwargs)

property tot_trans

f(out_dim=None, remove_constant=True, n_skip: int | None = None)

g(out_dim=None, 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) → tuple[list[jax.Array], list[jax.Array], list[jax.Array] | None, bool]

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

get_n_skip(n_skip: int | None = None, remove_constant=True)

timescales(remove_constant=True, n_skip=None)

IMLCV.base.rounds.X

IMLCV.base.rounds.X2

class IMLCV.base.rounds.Covariances(*args, **kwargs)

Bases: IMLCV.base.datastructures.MyPyTreeNode

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

rho_00: jax.Array

rho_01: jax.Array | None

rho_10: jax.Array | None

rho_11: jax.Array | None

pi_s_0: jax.Array | None

pi_s_1: jax.Array | None

sigma_0: jax.Array

sigma_1: jax.Array | None

W_0: jax.Array | None = None

W_1: jax.Array | None = None

only_diag: bool = False

trans_f: IMLCV.base.CV.CvTrans | IMLCV.base.CV.CvTrans | None = None

trans_g: IMLCV.base.CV.CvTrans | IMLCV.base.CV.CvTrans | None = None

symmetric: bool = False

static create(cv_0: list[X], cv_1: list[X] | None = None, nl: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, nl_t: list[IMLCV.base.CV.NeighbourList] | IMLCV.base.CV.NeighbourList | None = None, w: list[jax.Array] | None = None, w_t: list[jax.Array] | None = None, calc_pi=False, macro_chunk=1000, chunk_size=None, only_diag=False, trans_f: IMLCV.base.CV.CvTrans | IMLCV.base.CV.CvTrans | None = None, trans_g: IMLCV.base.CV.CvTrans | IMLCV.base.CV.CvTrans | None = None, symmetric=False, calc_C00=True, calc_C01=True, calc_C10=True, calc_C11=True, shmap_kwargs=ShmapKwargs.create(), verbose=True) → Covariances

property pi_0

property pi_1

property C00

property C01

property C10

property C11

property sigma_0_inv

property sigma_1_inv

whiten_rho(choice, epsilon: float = 0.0001, max_features=None, verbose=False) → jax.Array

mask(eps_pre: float | None, max_features: int = 2000, auto_cov_threshold: float | None = None)

decompose(out_dim: int | None = None, sparse=True, out_eps: float | None = None)

symmetrize()