Index _ | A | B | C | D | E | F | G | H | I | J | K | L | M | N | O | P | R | S | T | U | V | W _ __call__() (jaqmc.optimizer.base.OptimizerInit method) (jaqmc.optimizer.base.OptimizerUpdate method) (jaqmc.sampler.base.BatchLogProb method) (jaqmc.sampler.base.SamplerInit method) (jaqmc.sampler.base.SamplerStep method) (jaqmc.sampler.mcmc.SamplingProposal method) (jaqmc.wavefunction.Wavefunction method) (jaqmc.wavefunction.WavefunctionEvaluate method) (jaqmc.wavefunction.WavefunctionInit method) A abs_isotropic (jaqmc.wavefunction.output.envelope.EnvelopeType attribute) abs_logdets (jaqmc.wavefunction.output.logdet.ComplexLogDetOutput attribute) (jaqmc.wavefunction.output.logdet.RealLogDetOutput attribute) adam (class in jaqmc.optimizer.optax) addressable_data() (in module jaqmc.utils.parallel_jax) ae_features (jaqmc.wavefunction.input.atomic.AtomicEmbedding attribute) ae_vec (jaqmc.wavefunction.input.atomic.AtomicEmbedding attribute) aggregate_features() (jaqmc.wavefunction.backbone.ferminet.FermiLayers method) all_gather() (in module jaqmc.utils.parallel_jax) (jaqmc.data.BatchedData method) angstrom (jaqmc.utils.units.LengthUnit attribute) annotate_yaml_with_sources() (in module jaqmc.utils.yaml_format) array_partitions() (in module jaqmc.utils.array) ArrayLikeTree (in module jaqmc.array_types) ArrayTree (in module jaqmc.array_types) Atom (class in jaqmc.utils.atomic.atom) atom_list (jaqmc.utils.atomic.gto.AtomicOrbitalEvaluator attribute) atomic_number (jaqmc.utils.atomic.elements.Element attribute) AtomicEmbedding (class in jaqmc.wavefunction.input.atomic) AtomicOrbitalEvaluator (class in jaqmc.utils.atomic.gto) AtomicSystemConfig (class in jaqmc.utils.atomic.atomic_system) B basis (jaqmc.utils.atomic.scf.MolecularSCF attribute) basis_dict (jaqmc.utils.atomic.gto.AtomicOrbitalEvaluator attribute) batch_size (jaqmc.data.BatchedData property) BatchedData (class in jaqmc.data) BatchLogProb (class in jaqmc.sampler.base) bcc (jaqmc.geometry.pbc.SymmetryType attribute) bohr (jaqmc.utils.units.LengthUnit attribute) build() (jaqmc.workflow.stage.evaluation.EvalStageBuilder method) (jaqmc.workflow.stage.sampling.SamplingStageBuilder method) (jaqmc.workflow.stage.vmc.VMCStageBuilder method) build_distance_fn() (in module jaqmc.geometry.pbc) C cart2sph() (in module jaqmc.utils.atomic.gto) CartesianAxis (class in jaqmc.estimator.density) (class in jaqmc.estimator.density.cartesian) CartesianDensity (class in jaqmc.estimator.density) check() (jaqmc.data.BatchedData method) check_wired() (in module jaqmc.utils.wiring) CompatibleFunc (in module jaqmc.utils.func_transform) ComplexLogDetOutput (class in jaqmc.wavefunction.output.logdet) ComplexWFOutput (class in jaqmc.wavefunction.base) config_class (jaqmc.workflow.stage.sampling.SamplingStageBuilder attribute) ConfigCheck (class in jaqmc.workflow.base) ConfigManager (class in jaqmc.utils.config) ConfigManagerLike (class in jaqmc.utils.config) configurable_dataclass() (in module jaqmc.utils.config) configure_estimators() (jaqmc.workflow.stage.evaluation.EvalStageBuilder method) (jaqmc.workflow.stage.sampling.SamplingStageBuilder method) (jaqmc.workflow.stage.vmc.VMCStageBuilder method) configure_loss_grads() (jaqmc.workflow.stage.vmc.VMCStageBuilder method) configure_optimizer() (jaqmc.workflow.stage.vmc.VMCStageBuilder method) configure_sample_plan() (jaqmc.workflow.stage.evaluation.EvalStageBuilder method) (jaqmc.workflow.stage.sampling.SamplingStageBuilder method) (jaqmc.workflow.stage.vmc.VMCStageBuilder method) configure_writers() (jaqmc.workflow.stage.evaluation.EvalStageBuilder method) (jaqmc.workflow.stage.sampling.SamplingStageBuilder method) (jaqmc.workflow.stage.vmc.VMCStageBuilder method) ConsoleWriter (class in jaqmc.writer.console) Constant (class in jaqmc.optimizer.schedule) coulomb (jaqmc.app.hall.config.InteractionType attribute) counter (jaqmc.sampler.mcmc.MCMCState attribute) create_state() (jaqmc.workflow.stage.base.WorkStage method) CSVWriter (class in jaqmc.writer.csv) D Data (class in jaqmc.data) data (jaqmc.data.BatchedData attribute) data_init (jaqmc.workflow.evaluation.EvaluationWorkflow attribute) (jaqmc.workflow.vmc.VMCWorkflow attribute) diagonal (jaqmc.wavefunction.output.envelope.EnvelopeType attribute) digest() (jaqmc.estimator.base.EstimatorPipeline method) distance_type (jaqmc.wavefunction.input.atomic.SolidFeatures attribute) DistanceType (class in jaqmc.geometry.pbc) distribute_spins() (in module jaqmc.utils.atomic.initialization) DistributedConfig (class in jaqmc.utils.parallel_jax) dump_yaml() (in module jaqmc.utils.yaml_format) E ECPEnergy (class in jaqmc.estimator.ecp.estimator) ee_features (jaqmc.wavefunction.input.atomic.AtomicEmbedding attribute) Element (class in jaqmc.utils.atomic.elements) energy() (jaqmc.estimator.ewald.EwaldSum method) ensure_configured() (jaqmc.workflow.stage.sampling.SamplingStageBuilder method) Envelope (class in jaqmc.wavefunction.output.envelope) envelope_type (jaqmc.wavefunction.output.envelope.Envelope attribute) EnvelopeType (class in jaqmc.wavefunction.output.envelope) EstimateFn (in module jaqmc.estimator.base) Estimator (class in jaqmc.estimator.base) EstimatorLike (in module jaqmc.estimator.base) EstimatorPipeline (class in jaqmc.estimator.base) EuclideanKinetic (class in jaqmc.estimator.kinetic) eval_aos (jaqmc.utils.atomic.gto.PBCAtomicOrbitalEvaluator attribute) eval_mos() (jaqmc.utils.atomic.scf.MolecularSCF method) (jaqmc.utils.atomic.scf.PeriodicSCF method) eval_orbitals() (jaqmc.utils.atomic.scf.MolecularSCF method) (jaqmc.utils.atomic.scf.PeriodicSCF method) eval_slater() (jaqmc.utils.atomic.scf.MolecularSCF method) (jaqmc.utils.atomic.scf.PeriodicSCF method) eval_stage (jaqmc.workflow.evaluation.EvaluationWorkflow attribute) EvalStageBuilder (class in jaqmc.workflow.stage.evaluation) evaluate (jaqmc.wavefunction.WavefunctionLike attribute) evaluate() (jaqmc.estimator.base.EstimatorPipeline method) (jaqmc.wavefunction.Wavefunction method) evaluate_batch() (jaqmc.estimator.base.Estimator method) evaluate_local() (jaqmc.estimator.base.Estimator method) (jaqmc.estimator.base.FunctionEstimator method) EvaluationWorkflow (class in jaqmc.workflow.evaluation) EvaluationWorkflowConfig (class in jaqmc.workflow.evaluation) EvaluationWorkStage (class in jaqmc.workflow.stage.evaluation) EvaluationWorkStageConfig (class in jaqmc.workflow.stage.evaluation) EwaldSum (class in jaqmc.estimator.ewald) F fcc (jaqmc.geometry.pbc.SymmetryType attribute) FermiLayers (class in jaqmc.wavefunction.backbone.ferminet) field_names (jaqmc.data.Data property) fields_with_batch (jaqmc.data.BatchedData attribute) FieldSpec (class in jaqmc.writer.console) finalize() (jaqmc.utils.config.ConfigManager method) finalize_state() (jaqmc.estimator.base.Estimator method) finalize_stats() (jaqmc.estimator.base.Estimator method) (jaqmc.estimator.base.EstimatorPipeline method) fori_loop (jaqmc.estimator.kinetic.LaplacianMode attribute) forward_laplacian (jaqmc.estimator.kinetic.LaplacianMode attribute) FractionalAxis (class in jaqmc.estimator.density) (class in jaqmc.estimator.density.fractional) FractionalDensity (class in jaqmc.estimator.density) from_pyscf() (jaqmc.utils.atomic.gto.AtomicOrbitalEvaluator class method) fully_batched_data() (jaqmc.data.BatchedData method) FunctionEstimator (class in jaqmc.estimator.base) G get() (jaqmc.utils.config.ConfigManager method) (jaqmc.utils.config.ConfigManagerLike method) (jaqmc.utils.config.ScopedConfigManager method) get_collection() (jaqmc.utils.config.ConfigManager method) (jaqmc.utils.config.ConfigManagerLike method) (jaqmc.utils.config.ScopedConfigManager method) get_core_electrons() (in module jaqmc.utils.atomic.ecp) get_distance_function() (in module jaqmc.geometry.pbc) get_module() (jaqmc.utils.config.ConfigManager method) (jaqmc.utils.config.ConfigManagerLike method) (jaqmc.utils.config.ScopedConfigManager method) get_orbital_kpoints() (jaqmc.utils.atomic.scf.PeriodicSCF method) get_reciprocal_vectors() (in module jaqmc.utils.supercell) get_supercell_copies() (in module jaqmc.utils.supercell) get_supercell_kpts() (in module jaqmc.utils.supercell) get_symmetry_lat() (in module jaqmc.geometry.pbc) get_valence_spin_config() (in module jaqmc.utils.atomic.ecp) grad_maybe_complex() (in module jaqmc.utils.func_transform) H HDF5Writer (class in jaqmc.writer.hdf5) hessian_maybe_complex() (in module jaqmc.utils.func_transform) hexagonal (jaqmc.geometry.pbc.SymmetryType attribute) I import_module_or_file() (in module jaqmc.utils.module_resolver) init (jaqmc.optimizer.base.OptimizerLike attribute) (jaqmc.sampler.SamplerLike attribute) init() (jaqmc.estimator.base.Estimator method) (jaqmc.estimator.base.EstimatorPipeline method) (jaqmc.sampler.mcmc.MCMCSampler method) (jaqmc.sampler.SamplePlan method) init_params (jaqmc.wavefunction.WavefunctionLike attribute) init_params() (jaqmc.wavefunction.Wavefunction method) initialize_electrons_gaussian() (in module jaqmc.utils.atomic.initialization) InteractionType (class in jaqmc.app.hall.config) iqr_clip() (in module jaqmc.utils.clip) iqr_clip_real() (in module jaqmc.utils.clip) isotropic (jaqmc.wavefunction.output.envelope.EnvelopeType attribute) J JastrowType (class in jaqmc.app.molecule.wavefunction.psiformer) jit_sharded() (in module jaqmc.utils.parallel_jax) K KFACOptimizer (class in jaqmc.optimizer.kfac.kfac) L lamb (class in jaqmc.optimizer.optax) LaplacianMode (class in jaqmc.estimator.kinetic) lattice_vectors (jaqmc.utils.atomic.gto.PBCAtomicOrbitalEvaluator attribute) LayerNormMode (class in jaqmc.wavefunction.backbone.psiformer) LengthUnit (class in jaqmc.utils.units) linearize_maybe_complex() (in module jaqmc.utils.func_transform) LogDet (class in jaqmc.wavefunction.output.logdet) logpsi (jaqmc.wavefunction.output.logdet.ComplexLogDetOutput attribute) (jaqmc.wavefunction.output.logdet.RealLogDetOutput attribute) LogPsiWFOutput (class in jaqmc.wavefunction.base) loop() (jaqmc.workflow.stage.base.WorkStage method) LossAndGrad (class in jaqmc.estimator.loss_grad) M make_mesh() (in module jaqmc.utils.parallel_jax) make_pbc_gaussian_proposal() (in module jaqmc.geometry.pbc) make_pretrain_log_amplitude() (in module jaqmc.utils.atomic.pretrain) make_pretrain_loss() (in module jaqmc.utils.atomic.pretrain) make_sharding() (in module jaqmc.utils.parallel_jax) match_first_axis_of() (in module jaqmc.utils.array) MCMCSampler (class in jaqmc.sampler.mcmc) MCMCState (class in jaqmc.sampler.mcmc) mean_field (jaqmc.utils.atomic.scf.MolecularSCF attribute) merge() (jaqmc.data.Data method) minimal (jaqmc.geometry.pbc.SymmetryType attribute) mo_coeff (jaqmc.utils.atomic.scf.MolecularSCF attribute) MolecularSCF (class in jaqmc.utils.atomic.scf) molecule (jaqmc.utils.atomic.scf.MolecularSCF attribute) MoleculeFeatures (class in jaqmc.wavefunction.input.atomic) N name (jaqmc.utils.config.ConfigManagerLike property) ndets (jaqmc.wavefunction.output.envelope.Envelope attribute) nelectrons (jaqmc.utils.atomic.scf.MolecularSCF attribute) NONE (jaqmc.app.molecule.wavefunction.psiformer.JastrowType attribute) nspins (jaqmc.wavefunction.output.envelope.Envelope attribute) nu (jaqmc.geometry.pbc.DistanceType attribute) null (jaqmc.wavefunction.backbone.psiformer.LayerNormMode attribute) (jaqmc.wavefunction.output.envelope.EnvelopeType attribute) NumericWavefunctionEvaluate (in module jaqmc.wavefunction.base) NumPyCheckpointManager (class in jaqmc.utils.checkpoint) O open() (jaqmc.writer.base.Writer method) (jaqmc.writer.base.Writers method) (jaqmc.writer.console.ConsoleWriter method) (jaqmc.writer.csv.CSVWriter method) (jaqmc.writer.hdf5.HDF5Writer method) OptimizerInit (class in jaqmc.optimizer.base) OptimizerLike (class in jaqmc.optimizer.base) OptimizerUpdate (class in jaqmc.optimizer.base) OrbitalProjection (class in jaqmc.wavefunction.output.orbital) orbitals_spin_split (jaqmc.wavefunction.output.envelope.Envelope attribute) P pair_displacements_between() (in module jaqmc.geometry.obc) pair_displacements_within() (in module jaqmc.geometry.obc) Params (in module jaqmc.array_types) parse() (jaqmc.writer.console.FieldSpec static method) partition() (jaqmc.workflow.stage.sampling.SamplingState method) partition_spec (jaqmc.data.BatchedData property) PathLike (in module jaqmc.utils.checkpoint) PBCAtomicOrbitalEvaluator (class in jaqmc.utils.atomic.gto) period (jaqmc.utils.atomic.elements.Element attribute) PeriodicSCF (class in jaqmc.utils.atomic.scf) pmean() (in module jaqmc.utils.parallel_jax) pmoves (jaqmc.sampler.mcmc.MCMCState attribute) post (jaqmc.wavefunction.backbone.psiformer.LayerNormMode attribute) pre (jaqmc.wavefunction.backbone.psiformer.LayerNormMode attribute) prepare() (jaqmc.workflow.base.Workflow method) pretrain_stage (jaqmc.workflow.vmc.VMCWorkflow attribute) primitive_lattice (jaqmc.wavefunction.input.atomic.SolidFeatures attribute) PRNGKey (in module jaqmc.array_types) PsiformerBackbone (class in jaqmc.wavefunction.backbone.psiformer) PsiformerLayer (class in jaqmc.wavefunction.backbone.psiformer) pvary() (in module jaqmc.utils.parallel_jax) pyscf_mol (jaqmc.utils.atomic.scf.MolecularSCF attribute) PyTree (in module jaqmc.array_types) R r_ae (jaqmc.wavefunction.input.atomic.AtomicEmbedding attribute) RealLogDetOutput (class in jaqmc.wavefunction.output.logdet) RealWFOutput (class in jaqmc.wavefunction.base) reduce() (jaqmc.estimator.base.Estimator method) rescale (jaqmc.wavefunction.input.atomic.MoleculeFeatures attribute) resolve_object() (in module jaqmc.utils.module_resolver) resolve_path_template() (jaqmc.writer.base.Writer static method) restore() (jaqmc.utils.checkpoint.NumPyCheckpointManager method) restore_checkpoint() (jaqmc.workflow.stage.base.WorkStage method) (jaqmc.workflow.vmc.VMCWorkflow method) restore_from_file() (jaqmc.utils.checkpoint.NumPyCheckpointManager static method) restore_path (jaqmc.utils.checkpoint.NumPyCheckpointManager attribute) restricted (jaqmc.utils.atomic.scf.MolecularSCF attribute) run() (jaqmc.utils.atomic.scf.MolecularSCF method) (jaqmc.utils.atomic.scf.PeriodicSCF method) (jaqmc.workflow.base.Workflow method) (jaqmc.workflow.evaluation.EvaluationWorkflow method) (jaqmc.workflow.stage.base.WorkStage method) (jaqmc.workflow.vmc.VMCWorkflow method) RunContext (class in jaqmc.workflow.stage.base) runtime_dep() (in module jaqmc.utils.wiring) S SamplePlan (class in jaqmc.sampler) SamplerInit (class in jaqmc.sampler.base) SamplerLike (class in jaqmc.sampler) SamplerStep (class in jaqmc.sampler.base) SamplingProposal (class in jaqmc.sampler.mcmc) SamplingStageBuilder (class in jaqmc.workflow.stage.sampling) SamplingState (class in jaqmc.workflow.stage.sampling) save() (jaqmc.utils.checkpoint.NumPyCheckpointManager method) save_path (jaqmc.utils.checkpoint.NumPyCheckpointManager attribute) scaled_f() (in module jaqmc.geometry.pbc) scaled_g() (in module jaqmc.geometry.pbc) scan (jaqmc.estimator.kinetic.LaplacianMode attribute) scoped() (jaqmc.utils.config.ConfigManager method) ScopedConfigManager (class in jaqmc.utils.config) setup() (jaqmc.wavefunction.input.atomic.SolidFeatures method) (jaqmc.wavefunction.output.envelope.Envelope method) sign_logdets (jaqmc.wavefunction.output.logdet.ComplexLogDetOutput attribute) (jaqmc.wavefunction.output.logdet.RealLogDetOutput attribute) sign_logpsi (jaqmc.wavefunction.output.logdet.RealLogDetOutput attribute) SIMPLE_EE (jaqmc.app.molecule.wavefunction.psiformer.JastrowType attribute) SimpleEEJastrow (class in jaqmc.wavefunction.jastrow) simulation_lattice (jaqmc.wavefunction.input.atomic.SolidFeatures attribute) solid_harmonic() (in module jaqmc.utils.atomic.gto) SolidFeatures (class in jaqmc.wavefunction.input.atomic) spec (jaqmc.utils.atomic.gto.AtomicOrbitalEvaluator attribute) sphere_proposal() (in module jaqmc.geometry.sphere) SphericalDensity (class in jaqmc.estimator.density) SphericalKinetic (class in jaqmc.estimator.kinetic) spin (jaqmc.utils.atomic.elements.Element attribute) SpinSquared (class in jaqmc.estimator.spin) split_nonempty_channels() (in module jaqmc.utils.array) SplitChannelDense (class in jaqmc.wavefunction.output.orbital) SROptimizer (class in jaqmc.optimizer.sr) StageAbort (class in jaqmc.workflow.stage.base) Standard (class in jaqmc.optimizer.schedule) stddev (jaqmc.sampler.mcmc.MCMCState attribute) step (jaqmc.sampler.SamplerLike attribute) step() (jaqmc.sampler.mcmc.MCMCSampler method) (jaqmc.sampler.SamplePlan method) subset() (jaqmc.data.Data method) sym_type (jaqmc.wavefunction.input.atomic.SolidFeatures attribute) symbol (jaqmc.utils.atomic.elements.Element attribute) SymmetryType (class in jaqmc.geometry.pbc) T to_scalar() (jaqmc.writer.base.Writer static method) to_yaml() (jaqmc.utils.config.ConfigManager method) TotalEnergy (class in jaqmc.estimator.total_energy) train_stage (jaqmc.workflow.vmc.VMCWorkflow attribute) transform_maybe_complex() (in module jaqmc.utils.func_transform) transform_with_data() (in module jaqmc.utils.func_transform) tree_from_npz() (in module jaqmc.utils.checkpoint) tree_to_npz() (in module jaqmc.utils.checkpoint) tri (jaqmc.geometry.pbc.DistanceType attribute) U unbatched_example() (jaqmc.data.BatchedData method) unpaired_electron (jaqmc.utils.atomic.elements.Element property) update (jaqmc.optimizer.base.OptimizerLike attribute) use_preset() (jaqmc.utils.config.ConfigManager method) V vmap_axis (jaqmc.data.BatchedData property) VMCStageBuilder (class in jaqmc.workflow.stage.vmc) VMCState (class in jaqmc.workflow.stage.vmc) VMCWorkflow (class in jaqmc.workflow.vmc) VMCWorkStage (class in jaqmc.workflow.stage.vmc) VMCWorkStageConfig (class in jaqmc.workflow.stage.vmc) W Wavefunction (class in jaqmc.wavefunction) WavefunctionEvaluate (class in jaqmc.wavefunction) WavefunctionInit (class in jaqmc.wavefunction) WavefunctionLike (class in jaqmc.wavefunction) will_sample() (jaqmc.sampler.SamplePlan method) wire() (in module jaqmc.utils.wiring) with_imag() (in module jaqmc.utils.func_transform) with_output() (in module jaqmc.utils.func_transform) with_real() (in module jaqmc.utils.func_transform) Workflow (class in jaqmc.workflow.base) WorkflowConfig (class in jaqmc.workflow.base) WorkStage (class in jaqmc.workflow.stage.base) WorkStageConfig (class in jaqmc.workflow.stage.base) wrap_positions() (in module jaqmc.geometry.pbc) write() (jaqmc.writer.base.Writer method) (jaqmc.writer.base.Writers method) (jaqmc.writer.console.ConsoleWriter method) (jaqmc.writer.csv.CSVWriter method) (jaqmc.writer.hdf5.HDF5Writer method) Writer (class in jaqmc.writer.base) Writers (class in jaqmc.writer.base)
