Modules

• Command Line Interface
  • docs()
  • eval()
  • md()
  • schema()
  • template_md_config()
  • template_train_config()
  • train()
  • validate_md_config()
  • validate_train_config()
  • version_callback()
  • visualize_model()
• Configuration
  • Training Configuration
  • Molecular Dynamics Configuration
• Data Pipeline
  • compute_nl()
  • get_shrink_wrapped_cell()
  • prefetch_to_single_device()
  • load_data_files()
  • find_largest_system()
  • pad_nl()
  • DatasetStats
• Optimizers
  • get_opt()
  • get_schedule()
  • map_nested_fn()
• Layers
  • EmpiricalEnergyTerm
  • ZBLRepulsion
  • uniform_range()
  • NTKLinear
  • stress_times_vol()
  • AtomisticReadout
  • PerElementScaleShift
• Model
  • AtomisticModel
  • EnergyDerivativeModel
  • EnergyModel
• Training
  • CSVLoggerApax
  • canonicalize_energy_grad_model_parameters()
  • canonicalize_energy_model_parameters()
  • check_for_ensemble()
  • restore_parameters()
  • restore_single_parameters()
  • stack_parameters()
  • eval_model()
  • load_test_data()
  • predict()
  • Loss
  • LossCollection
  • force_angle_div_force_label()
  • force_angle_exponential_weight()
  • force_angle_loss()
  • weighted_huber_loss()
  • weighted_squared_error()
  • Averagefp64
  • RootAverage
  • cosine_sim()
  • initialize_metrics()
  • mae_fn()
  • make_single_metric()
  • mse_fn()
  • initialize_datasets()
  • initialize_loss_fn()
  • run()
  • setup_logging()
  • fit()
  • global_norm()
• Molecular Dynamics
  • ASECalculator
  • FunctionTransformation
  • GaussianAcceleratedMolecularDynamics
  • UncertaintyDrivenDynamics
  • md_setup()
  • run_md()
  • run_nvt()
• Transfer Learning
• Batch Active Learning
  • compute_features()
  • feature_fn
  • dataset
  • create_feature_fn()
  • model
  • params
  • base_feature_map
  • is_ensemble
  • kernel_selection()
  • model_dir
  • train_atoms
  • pool_atoms
  • base_fm_options
  • selection_method
  • feature_transforms
  • selection_batch_size
  • processing_batch_size
  • FeatureTransformation
  • IdentityFeatures
  • LastLayerGradientFeatures
  • extract_feature_params()
  • KernelMatrix
  • max_dist_selection()
  • matrix
  • batch_size
  • batch_features()
  • ensemble_features()
• Utils
  • atoms_to_inputs()
  • atoms_to_labels()
  • tf_to_jax_dict()
  • load_data()
  • split_atoms()