Index

A

apax.bal.api module apax.bal.feature_maps module apax.bal.kernel module apax.bal.selection module apax.bal.transforms module apax.cli.apax_app module apax.data.initialization module apax.data.input_pipeline module apax.data.preprocessing module apax.data.statistics module apax.layers.activation module apax.layers.empirical module apax.layers.initializers module apax.layers.masking module apax.layers.ntk_linear module apax.layers.properties module apax.layers.readout module apax.layers.scaling module apax.md.ase_calc module apax.md.function_transformations module

apax.md.io module apax.md.nvt module apax.model.builder module apax.model.gmnn module apax.optimizer.get_optimizer module apax.train.callbacks module apax.train.checkpoints module apax.train.eval module apax.train.loss module apax.train.metrics module apax.train.run module apax.train.trainer module apax.transfer_learning.parameter_transfer module apax.utils.convert module apax.utils.data module apax.utils.math module apax.utils.random module ASECalculator (class in apax.md.ase_calc) AtomisticModel (class in apax.model.gmnn) AtomisticReadout (class in apax.layers.readout) atoms_to_inputs() (in module apax.utils.convert) atoms_to_labels() (in module apax.utils.convert) Averagefp64 (class in apax.train.metrics)

B

base_feature_map (in module apax.bal.api) base_fm_options (in module apax.bal.api)

batch_eval() (apax.md.ase_calc.ASECalculator method) batch_features() (in module apax.bal.transforms) batch_size (in module apax.bal.selection)

C

calculate() (apax.md.ase_calc.ASECalculator method) canonicalize_energy_grad_model_parameters() (in module apax.train.checkpoints) canonicalize_energy_model_parameters() (in module apax.train.checkpoints) check_for_ensemble() (in module apax.train.checkpoints) CheckpointConfig (class in apax.config.train_config) compute() (apax.train.metrics.RootAverage method)

compute_features() (in module apax.bal.api) compute_nl() (in module apax.data.preprocessing) Config (class in apax.config.train_config) cosine_sim() (in module apax.train.metrics) create_feature_fn() (in module apax.bal.api) CSVCallback (class in apax.config.train_config) CSVLoggerApax (class in apax.train.callbacks)

D

DataConfig (class in apax.config.train_config) dataset (in module apax.bal.api) DatasetStats (class in apax.data.statistics)

docs() (in module apax.cli.apax_app) dump_config() (apax.config.md_config.MDConfig method) (apax.config.train_config.Config method)

E

EmpiricalEnergyTerm (class in apax.layers.empirical) empty() (apax.train.metrics.Averagefp64 class method) EnergyDerivativeModel (class in apax.model.gmnn) EnergyModel (class in apax.model.gmnn)

ensemble_features() (in module apax.bal.transforms) eval() (in module apax.cli.apax_app) eval_model() (in module apax.train.eval) extract_feature_params() (in module apax.bal.feature_maps)

F

feature_fn (in module apax.bal.api) feature_transforms (in module apax.bal.api) FeatureTransformation (class in apax.bal.feature_maps) find_largest_system() (in module apax.data.input_pipeline)

fit() (in module apax.train.trainer) force_angle_div_force_label() (in module apax.train.loss) force_angle_exponential_weight() (in module apax.train.loss) force_angle_loss() (in module apax.train.loss) FunctionTransformation (class in apax.md.function_transformations)

G

GaussianAcceleratedMolecularDynamics (class in apax.md.function_transformations) get_opt() (in module apax.optimizer.get_optimizer)

get_schedule() (in module apax.optimizer.get_optimizer) get_shrink_wrapped_cell() (in module apax.data.preprocessing) global_norm() (in module apax.train.trainer)

I

IdentityFeatures (class in apax.bal.feature_maps) implemented_properties (apax.md.ase_calc.ASECalculator attribute) initialize_datasets() (in module apax.train.run)

initialize_loss_fn() (in module apax.train.run) initialize_metrics() (in module apax.train.metrics) is_ensemble (in module apax.bal.api)

K

kernel_selection() (in module apax.bal.api)

KernelMatrix (class in apax.bal.kernel)

L

LastLayerGradientFeatures (class in apax.bal.feature_maps) load_data() (in module apax.utils.data) load_data_files() (in module apax.data.initialization)

load_test_data() (in module apax.train.eval) Loss (class in apax.train.loss) LossCollection (class in apax.train.loss) LossConfig (class in apax.config.train_config)

M

mae_fn() (in module apax.train.metrics) make_single_metric() (in module apax.train.metrics) map_nested_fn() (in module apax.optimizer.get_optimizer) matrix (in module apax.bal.selection) max_dist_selection() (in module apax.bal.selection) md() (in module apax.cli.apax_app) md_setup() (in module apax.md.nvt) MDConfig (class in apax.config.md_config) MetricsConfig (class in apax.config.train_config) MLFlowCallback (class in apax.config.train_config) model (in module apax.bal.api) model_computed_fields (apax.bal.feature_maps.FeatureTransformation attribute) (apax.bal.feature_maps.IdentityFeatures attribute) (apax.bal.feature_maps.LastLayerGradientFeatures attribute) model_config (apax.bal.feature_maps.FeatureTransformation attribute) (apax.bal.feature_maps.IdentityFeatures attribute) (apax.bal.feature_maps.LastLayerGradientFeatures attribute) model_dir (in module apax.bal.api) model_fields (apax.bal.feature_maps.FeatureTransformation attribute) (apax.bal.feature_maps.IdentityFeatures attribute) (apax.bal.feature_maps.LastLayerGradientFeatures attribute) ModelConfig (class in apax.config.train_config) module apax.bal.api apax.bal.feature_maps apax.bal.kernel apax.bal.selection apax.bal.transforms apax.cli.apax_app apax.data.initialization apax.data.input_pipeline apax.data.preprocessing apax.data.statistics apax.layers.activation apax.layers.empirical apax.layers.initializers apax.layers.masking apax.layers.ntk_linear apax.layers.properties apax.layers.readout apax.layers.scaling apax.md.ase_calc apax.md.function_transformations apax.md.io apax.md.nvt apax.model.builder apax.model.gmnn apax.optimizer.get_optimizer apax.train.callbacks apax.train.checkpoints apax.train.eval apax.train.loss apax.train.metrics apax.train.run apax.train.trainer apax.transfer_learning.parameter_transfer apax.utils.convert apax.utils.data apax.utils.math apax.utils.random

mse_fn() (in module apax.train.metrics)

N

NPTOptions (class in apax.config.md_config)

NTKLinear (class in apax.layers.ntk_linear)

O

on_test_batch_end() (apax.train.callbacks.CSVLoggerApax method)

OptimizerConfig (class in apax.config.train_config)

P

pad_nl() (in module apax.data.input_pipeline) params (in module apax.bal.api) PerElementScaleShift (class in apax.layers.scaling)

pool_atoms (in module apax.bal.api) predict() (in module apax.train.eval) prefetch_to_single_device() (in module apax.data.preprocessing) processing_batch_size (in module apax.bal.api)

R

restore_parameters() (in module apax.train.checkpoints) restore_single_parameters() (in module apax.train.checkpoints) RootAverage (class in apax.train.metrics)

run() (in module apax.train.run) run_md() (in module apax.md.nvt) run_nvt() (in module apax.md.nvt)

S

schema() (in module apax.cli.apax_app) score() (apax.bal.kernel.KernelMatrix method) selection_batch_size (in module apax.bal.api) selection_method (in module apax.bal.api) setup() (apax.layers.empirical.ZBLRepulsion method) (apax.layers.readout.AtomisticReadout method) (apax.layers.scaling.PerElementScaleShift method) (apax.model.gmnn.EnergyModel method)

setup_logging() (in module apax.train.run) split_atoms() (in module apax.utils.data) stack_parameters() (in module apax.train.checkpoints) stress_times_vol() (in module apax.layers.properties)

T

TBCallback (class in apax.config.train_config) template_md_config() (in module apax.cli.apax_app) template_train_config() (in module apax.cli.apax_app)

tf_to_jax_dict() (in module apax.utils.convert) train() (in module apax.cli.apax_app) train_atoms (in module apax.bal.api) TrainProgressbarConfig (class in apax.config.train_config)

U

UncertaintyDrivenDynamics (class in apax.md.function_transformations)

uniform_range() (in module apax.layers.initializers)

V

validate_md_config() (in module apax.cli.apax_app) validate_train_config() (in module apax.cli.apax_app)

version_callback() (in module apax.cli.apax_app) visualize_model() (in module apax.cli.apax_app)

W

weighted_huber_loss() (in module apax.train.loss)

weighted_squared_error() (in module apax.train.loss)

Z

ZBLRepulsion (class in apax.layers.empirical)