Modules

• Command Line Interface
  • docs()
  • eval()
  • md()
  • schema_md()
  • schema_train()
  • schema_vscode()
  • template_md_config()
  • template_train_config()
  • train()
  • validate_md_config()
  • validate_train_config()
  • version_callback()
  • visualize_model()
• Configuration
  • Training Configuration
  • Model Configuration
  • Molecular Dynamics Configuration
  • Hyperparameter Optimization Configuration
• Data Pipeline
  • compute_nl()
  • get_shrink_wrapped_cell()
  • prefetch_to_single_device()
  • load_data_files()
  • CachedInMemoryDataset
  • InMemoryDataset
  • OTFInMemoryDataset
  • PerBatchPaddedDataset
  • find_largest_system()
  • pad_nl()
  • round_up_to_multiple()
  • DatasetStats
• Optimizers
  • cyclic_cosine_decay_schedule()
  • get_opt()
  • get_schedule()
• Layers
  • get_activation_fn()
  • EmpiricalEnergyTerm
  • ExponentialRepulsion
  • LatentEwald
  • ZBLRepulsion
  • uniform_range()
  • NTKLinear
  • PropertyHead
  • stress_times_vol()
  • AtomisticReadout
  • PerElementScaleShift
• Model
  • EnergyDerivativeModel
  • EnergyModel
  • FeatureModel
  • ShallowEnsembleModel
• Training
  • CSVLoggerApax
  • KerasPruningCallback
  • 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
  • crps_loss()
  • force_angle_div_force_label()
  • force_angle_exponential_weight()
  • force_angle_loss()
  • nll_loss()
  • weighted_huber_loss()
  • weighted_squared_error()
  • Averagefp64
  • RootAverage
  • cosine_sim()
  • initialize_metrics()
  • mae_fn()
  • make_single_metric()
  • mse_fn()
  • per_atom_mae_fn()
  • per_atom_mse_fn()
  • initialize_datasets()
  • initialize_loss_fn()
  • run()
  • setup_logging()
  • EarlyStop
  • fit()
  • make_step_fns()
• Molecular Dynamics
  • ASECalculator
  • GaussianAcceleratedMolecularDynamics
  • GlobalCalibration
  • ProcessStress
  • UncertaintyDrivenDynamics
  • md_setup()
  • run_md()
  • run_sim()
  • DynamicsCheckBase
  • EnergyUncertaintyCheck
  • ForceUncertaintyCheck
  • RadiusCheck
  • ReflectionCheck
  • ConstraintBase
  • FixAtoms
  • FixCenterOfMass
  • FixRotation
  • BiasEnergyBase
  • SphericalWall
  • SimulationFunctions
  • System
• Transfer Learning
  • black_list_param_transfer()
  • transfer_parameters()
• Batch Active Learning
  • compute_features()
  • create_feature_fn()
  • kernel_selection()
  • FeatureTransformation
  • FullGradientRPFeatures
  • IdentityFeatures
  • LastLayerForceFeatures
  • LastLayerGradientFeatures
  • extract_feature_params()
  • KernelMatrix
  • max_dist_selection()
  • batch_features()
  • ensemble_features()
• Utils
  • atoms_to_inputs()
  • atoms_to_labels()
  • tf_to_jax_dict()
  • load_data()
  • split_atoms()
  • center_of_mass()
• Nodes