Index

A

activation_fn() (apax.layers.readout.AtomisticReadout method) 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.bias module apax.md.constraints module apax.md.dynamics_checks module apax.md.function_transformations module

apax.md.io module apax.md.md_checkpoint module apax.md.schedules module apax.md.sim_utils module apax.md.simulate module apax.nn.builder module apax.nn.models 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) atom_padding (apax.data.input_pipeline.PerBatchPaddedDataset attribute) 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) BesselBasisConfig (class in apax.config.model_config) BiasEnergyBase (class in apax.md.bias) black_list_param_transfer() (in module apax.transfer_learning.parameter_transfer)

C

CachedDataset (class in apax.config.train_config) CachedInMemoryDataset (class in apax.data.input_pipeline) 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) center_of_mass() (in module apax.utils.math) check_for_ensemble() (in module apax.train.checkpoints) 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) ConstraintBase (class in apax.md.constraints) Correction (class in apax.config.model_config) cosine_sim() (in module apax.train.metrics) create_feature_fn() (in module apax.bal.api) crps_loss() (in module apax.train.loss) CSVCallback (class in apax.config.train_config) CSVLoggerApax (class in apax.train.callbacks) cyclic_cosine_decay_schedule() (in module apax.optimizer.get_optimizer)

D

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

dtype (apax.layers.empirical.EmpiricalEnergyTerm attribute) (apax.layers.readout.AtomisticReadout attribute) (apax.layers.scaling.PerElementScaleShift attribute) dump_config() (apax.config.md_config.MDConfig method) (apax.config.train_config.Config method) DynamicsCheckBase (class in apax.md.dynamics_checks)

E

EarlyStop EmpiricalEnergyTerm (class in apax.layers.empirical) empty() (apax.train.metrics.Averagefp64 class method) EnergyDerivativeModel (class in apax.nn.models) EnergyModel (class in apax.nn.models) EnergyUncertaintyCheck (class in apax.md.dynamics_checks) enqueue_batches() (apax.data.input_pipeline.PerBatchPaddedDataset method)

ensemble_features() (in module apax.bal.transforms) EquivMPConfig (class in apax.config.model_config) eval() (in module apax.cli.apax_app) eval_model() (in module apax.train.eval) ExponentialRepulsion (class in apax.config.model_config) (class in apax.layers.empirical) extract_feature_params() (in module apax.bal.feature_maps)

F

feature_fn (in module apax.bal.api) feature_transforms (in module apax.bal.api) FeatureModel (class in apax.nn.models) FeatureTransformation (class in apax.bal.feature_maps) find_largest_system() (in module apax.data.input_pipeline) fit() (in module apax.train.trainer) FixAtoms (class in apax.md.constraints)

FixCenterOfMass (class in apax.md.constraints) FixRotation (class in apax.md.constraints) 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) ForceUncertaintyCheck (class in apax.md.dynamics_checks) FullEnsembleConfig (class in apax.config.model_config) FullGradientRPFeatures (class in apax.bal.feature_maps)

G

GaussianAcceleratedMolecularDynamics (class in apax.md.function_transformations) GaussianBasisConfig (class in apax.config.model_config) get_activation_fn() (in module apax.layers.activation) get_descriptors() (apax.md.ase_calc.ASECalculator method)

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) GlobalCalibration (class in apax.md.function_transformations) GMNNConfig (class in apax.config.model_config)

I

IdentityFeatures (class in apax.bal.feature_maps) implemented_properties (apax.md.ase_calc.ASECalculator attribute) init_input() (apax.data.input_pipeline.InMemoryDataset method) initialize_datasets() (in module apax.train.run)

initialize_loss_fn() (in module apax.train.run) initialize_metrics() (in module apax.train.metrics) InMemoryDataset (class in apax.data.input_pipeline) interval (apax.config.optuna_config.OptunaPrunerConfig attribute) is_ensemble (in module apax.bal.api)

K

KerasPruningCallback (class in apax.config.train_config) (class in apax.train.callbacks) kernel_selection() (in module apax.bal.api)

KernelMatrix (class in apax.bal.kernel) kwargs (apax.config.optuna_config.OptunaPrunerConfig attribute) (apax.config.optuna_config.OptunaSamplerConfig attribute)

L

LastLayerForceFeatures (class in apax.bal.feature_maps) LastLayerGradientFeatures (class in apax.bal.feature_maps) LatentEwald (class in apax.config.model_config) (class in apax.layers.empirical) 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) make_step_fns() (in module apax.train.trainer) 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.simulate) 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_config (apax.bal.feature_maps.FeatureTransformation attribute) (apax.bal.feature_maps.FullGradientRPFeatures attribute) (apax.bal.feature_maps.IdentityFeatures attribute) (apax.bal.feature_maps.LastLayerForceFeatures attribute) (apax.bal.feature_maps.LastLayerGradientFeatures attribute) (apax.md.bias.BiasEnergyBase attribute) (apax.md.bias.SphericalWall attribute) (apax.md.constraints.ConstraintBase attribute) (apax.md.constraints.FixAtoms attribute) (apax.md.constraints.FixCenterOfMass attribute) (apax.md.constraints.FixRotation attribute) (apax.md.dynamics_checks.DynamicsCheckBase attribute) (apax.md.dynamics_checks.EnergyUncertaintyCheck attribute) (apax.md.dynamics_checks.ForceUncertaintyCheck attribute) (apax.md.dynamics_checks.RadiusCheck attribute) (apax.md.dynamics_checks.ReflectionCheck attribute) model_dir (in module apax.bal.api) 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.bias apax.md.constraints apax.md.dynamics_checks apax.md.function_transformations apax.md.io apax.md.md_checkpoint apax.md.schedules apax.md.sim_utils apax.md.simulate apax.nn.builder apax.nn.models 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

monitor (apax.config.optuna_config.OptunaConfig attribute) mse_fn() (in module apax.train.metrics)

N

n_trials (apax.config.optuna_config.OptunaConfig attribute) name (apax.config.optuna_config.OptunaPrunerConfig attribute) (apax.config.optuna_config.OptunaSamplerConfig attribute) nl_padding (apax.data.input_pipeline.PerBatchPaddedDataset attribute)

nll_loss() (in module apax.train.loss) NPTOptions (class in apax.config.md_config) NTKLinear (class in apax.layers.ntk_linear) num_workers (apax.data.input_pipeline.PerBatchPaddedDataset attribute)

O

on_epoch_end() (apax.train.callbacks.CSVLoggerApax method) (apax.train.callbacks.KerasPruningCallback method) on_test_batch_end() (apax.train.callbacks.CSVLoggerApax method) OptimizerConfig (class in apax.config.train_config)

OptunaConfig (class in apax.config.optuna_config) OptunaPrunerConfig (class in apax.config.optuna_config) OptunaSamplerConfig (class in apax.config.optuna_config) OTFDataset (class in apax.config.train_config) OTFInMemoryDataset (class in apax.data.input_pipeline)

P

pad_nl() (in module apax.data.input_pipeline) params (in module apax.bal.api) PBPDatset (class in apax.config.train_config) per_atom_mae_fn() (in module apax.train.metrics) per_atom_mse_fn() (in module apax.train.metrics) PerBatchPaddedDataset (class in apax.data.input_pipeline) 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) ProcessStress (class in apax.md.function_transformations) PropertyHead (class in apax.config.model_config) (class in apax.layers.properties) pruner_config (apax.config.optuna_config.OptunaConfig attribute)

R

RadiusCheck (class in apax.md.dynamics_checks) ReflectionCheck (class in apax.md.dynamics_checks) restore_parameters() (in module apax.train.checkpoints) restore_single_parameters() (in module apax.train.checkpoints)

RootAverage (class in apax.train.metrics) round_up_to_multiple() (in module apax.data.input_pipeline) run() (in module apax.train.run) run_md() (in module apax.md.simulate) run_sim() (in module apax.md.simulate)

S

sampler_config (apax.config.optuna_config.OptunaConfig attribute) schema_md() (in module apax.cli.apax_app) schema_train() (in module apax.cli.apax_app) schema_vscode() (in module apax.cli.apax_app) score() (apax.bal.kernel.KernelMatrix method) search_space (apax.config.optuna_config.OptunaConfig attribute) seed (apax.config.optuna_config.OptunaConfig attribute) selection_batch_size (in module apax.bal.api) selection_method (in module apax.bal.api) setup() (apax.layers.empirical.ExponentialRepulsion method) (apax.layers.empirical.ZBLRepulsion method) (apax.layers.properties.PropertyHead method) (apax.layers.readout.AtomisticReadout method) (apax.layers.scaling.PerElementScaleShift method) (apax.nn.models.EnergyModel method) (apax.nn.models.FeatureModel method)

setup_logging() (in module apax.train.run) ShallowEnsembleConfig (class in apax.config.model_config) ShallowEnsembleModel (class in apax.nn.models) shuffle_and_batch() (apax.data.input_pipeline.CachedInMemoryDataset method) (apax.data.input_pipeline.OTFInMemoryDataset method) SimulationFunctions (class in apax.md.sim_utils) So3kratesConfig (class in apax.config.model_config) SphericalWall (class in apax.md.bias) split_atoms() (in module apax.utils.data) stack_parameters() (in module apax.train.checkpoints) steps_per_epoch() (apax.data.input_pipeline.InMemoryDataset method) stress_times_vol() (in module apax.layers.properties) study_log_file (apax.config.optuna_config.OptunaConfig attribute) study_name (apax.config.optuna_config.OptunaConfig attribute) System (class in apax.md.sim_utils)

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) transfer_parameters() (in module apax.transfer_learning.parameter_transfer) TransferLearningConfig (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

WeightAverage (class in apax.config.train_config)

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

Z

ZBLRepulsion (class in apax.config.model_config) (class in apax.layers.empirical)