refactor HF reference config and pretrain defaults

docs/api-reference/atomic.md

@@ -53,6 +53,8 @@ Atom definitions, element properties, electronic structure (SCF), orbital evalua
 ## Pretraining
 
 ```{eval-rst}
+.. autoclass:: jaqmc.utils.atomic.pretrain.PretrainReferenceConfig
+
 .. autofunction:: jaqmc.utils.atomic.pretrain.make_pretrain_log_amplitude
 .. autofunction:: jaqmc.utils.atomic.pretrain.make_pretrain_loss
 ```

docs/systems/molecule/index.md

@@ -5,7 +5,8 @@ boundary conditions. Most runs start from a YAML definition and a
 single `jaqmc molecule train` command. JaQMC then follows the standard
 molecular workflow:
 
-1. **Hartree-Fock (HF)** computes reference orbitals with PySCF.
+1. **Hartree-Fock (HF)** computes a reference electronic-structure solution with
+   PySCF.
 2. **Pretraining** matches the neural wavefunction to those orbitals.
 3. **VMC training** performs the main energy optimization.
 
@@ -76,14 +77,13 @@ that generate the underlying configuration for you.
 ### Single Atoms
 
 For a single atom, `system.module=atom` is a shortcut. You provide the element
-symbol and optional HF settings, and JaQMC fills in the matching electron spin
-configuration automatically.
+symbol, and JaQMC fills in the matching electron spin configuration
+automatically.
 
 ```yaml
 system:
   module: atom
   symbol: Li         # Element symbol (H, He, Li, Be, ...)
-  basis: sto-3g      # Basis set for SCF initialization
   # ecp: ccecp       # Optional: effective core potential
 ```
 
@@ -106,7 +106,6 @@ system:
   bond_length: 3.015  # Distance between atoms
   unit: bohr          # Length unit for bond_length
   spin: 0             # n_up - n_down for the full molecule
-  basis: cc-pvdz
 ```
 
 Save as `li_h_diatomic.yml`, then run:
@@ -118,33 +117,57 @@ jaqmc molecule train --yml li_h_diatomic.yml workflow.save_path=./runs/li_h_diat
 (molecule-basis-sets-and-ecps)=
 ## Basis Sets and ECPs
 
-The `basis` parameter controls the basis set used for the HF calculation. Any basis set supported by PySCF works:
+The examples above keep the Hartree-Fock reference at its default settings so
+you can focus on defining the system. When you do need to change the reference
+calculation itself, add a `pretrain.reference` section. That is where you set
+the HF basis and method. The `system` section still describes the physical
+system, including `system.ecp`.
+
+For example, to use a different HF basis:
+
+```yaml
+pretrain:
+  reference:
+    basis: cc-pvdz
+```
+
+Any basis set supported by PySCF works:
 
 - Minimal: `sto-3g` (default, fast)
 - Split-valence: `6-31g`, `6-311g`
 - Correlation-consistent: `cc-pvdz`, `cc-pvtz`, `cc-pvqz`
 
-For heavy elements (transition metals, lanthanides), use an effective core potential (ECP) to replace core electrons with a pseudopotential, reducing the number of electrons treated explicitly:
+For heavy elements (transition metals, lanthanides), use an effective core
+potential (ECP) to replace core electrons with a pseudopotential, reducing the
+number of electrons treated explicitly:
 
 ```yaml
 system:
   module: atom
   symbol: Fe
-  basis: ccecpccpvdz
   ecp: ccecp
+pretrain:
+  reference:
+    basis: ccecpccpvdz
 ```
 
 Both `basis` and `ecp` can be specified per element:
 
 ```yaml
+pretrain:
+  reference:
+    basis:
+      Fe: ccecpccpvdz
+      O: cc-pvdz
 system:
-  basis:
-    Fe: ccecpccpvdz
-    O: cc-pvdz
   ecp:
     Fe: ccecp
 ```
 
+If you need more control, the same `pretrain.reference` section also lets you
+change the HF method and pass extra PySCF options. See <project:train.md> for
+the full reference.
+
 ## Estimators
 
 The training stage computes energy from several components: kinetic energy,

docs/systems/molecule/train.md

@@ -185,6 +185,16 @@ Initializes the neural network to approximate Hartree-Fock orbitals before VMC
 training. It uses the same run, sampler, and writer schemas as the train stage,
 but with a different optimizer default and a workflow-wired supervised loss.
 
+### Reference (`pretrain.reference.*`)
+
+The Hartree-Fock reference is the PySCF calculation JaQMC uses to generate the
+target orbitals for pretraining. Most runs can keep the default settings.
+
+```{eval-rst}
+.. config-defaults:: jaqmc.app.molecule.config.base.MoleculePretrainReferenceConfig
+   :prefix: pretrain.reference
+```
+
 ### Run options (`pretrain.run.*`)
 
 ```{eval-rst}

docs/systems/solid/eval.md

@@ -31,6 +31,17 @@ are identical to the [training system config](#solid-train-system).
 Must match the training run. The effective defaults and built-in module choices
 are identical to the [training wavefunction config](#solid-train-wf).
 
+## Reference (`reference.*`)
+
+The Hartree-Fock reference is the PySCF calculation JaQMC uses when it needs
+reference orbitals or related setup from that calculation. Itis recommended to
+set the values to match the reference configuration used during training.
+
+```{eval-rst}
+.. config-defaults:: jaqmc.app.solid.config.base.SolidPretrainReferenceConfig
+   :prefix: reference
+```
+
 ## Run Options (`run.*`)
 
 Evaluation reuses the same checkpointing and sampling controls as training, but

docs/systems/solid/index.md

@@ -5,7 +5,8 @@ boundary conditions. Most runs start from a YAML definition and a
 single `jaqmc solid train` command. JaQMC then follows the same three-stage
 workflow used for [molecules](../molecule/index.md):
 
-1. **Hartree-Fock (HF)** computes reference orbitals with PySCF.
+1. **Hartree-Fock (HF)** computes a reference electronic-structure solution with
+   PySCF.
 2. **Pretraining** matches the neural wavefunction to those orbitals.
 3. **VMC training** performs the main energy optimization.
 
@@ -47,7 +48,6 @@ system:
     - symbol: H
       coords: [3.78, 3.78, 3.78]
   electron_spins: [2, 2]     # [n_up, n_down] per primitive cell
-  basis: sto-3g
 ```
 
 Then run training:
@@ -88,7 +88,6 @@ system:
   lattice_constant: 4.0     # in angstrom by default
   unit: angstrom            # or "bohr"
   # supercell: [2, 2, 2]    # Optional diagonal supercell shorthand
-  basis: sto-3g
 ```
 
 Save as `rock_salt.yml`, then run:
@@ -111,7 +110,6 @@ system:
   unit: bohr                 # or "angstrom"
   spin: 0                    # n_up - n_down per primitive cell
   # supercell: 4             # Optional repetition along the chain direction
-  basis: sto-3g
 ```
 
 Save as `two_atom_chain.yml`, then run:
@@ -120,8 +118,18 @@ Save as `two_atom_chain.yml`, then run:
 jaqmc solid train --yml two_atom_chain.yml workflow.save_path=./runs/two_atom_chain
 ```
 
+The examples above also keep the Hartree-Fock reference at its default
+settings. If you need a different HF basis or method, add a
+`pretrain.reference` section, for example:
+
+```yaml
+pretrain:
+  reference:
+    basis: cc-pvdz
+```
+
 Basis sets and ECPs work the same as for
-[molecules](#molecule-basis-sets-and-ecps).
+<project:#molecule-basis-sets-and-ecps>.
 
 ## Supercell Expansion
 

docs/systems/solid/train.md

@@ -172,6 +172,16 @@ Initializes the neural network to approximate Hartree-Fock orbitals before VMC
 training. It uses the same run, sampler, and writer schemas as the train stage,
 but with a different optimizer default and a workflow-wired supervised loss.
 
+### Reference (`pretrain.reference.*`)
+
+The Hartree-Fock reference is the PySCF calculation JaQMC uses to generate the
+target orbitals for pretraining. Most runs can keep the default settings.
+
+```{eval-rst}
+.. config-defaults:: jaqmc.app.solid.config.base.SolidPretrainReferenceConfig
+   :prefix: pretrain.reference
+```
+
 ### Run options (`pretrain.run.*`)
 
 ```{eval-rst}

src/jaqmc/app/molecule/config/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2025-2026 Bytedance Ltd. and/or its affiliates
 # SPDX-License-Identifier: Apache-2.0
 
-from .base import MoleculeConfig
+from .base import MoleculeConfig, MoleculePretrainReferenceConfig
 
-__all__ = ["MoleculeConfig"]
+__all__ = ["MoleculeConfig", "MoleculePretrainReferenceConfig"]

src/jaqmc/app/molecule/config/atom.py

@@ -11,15 +11,13 @@
 def atom_config(
     symbol: str = "H",
     electron_init_width: float = 1.0,
-    basis: str = "sto-3g",
     ecp: str | None = None,
 ):
     """Create a MoleculeConfig for a single atom.
 
     Args:
         symbol: Element symbol (e.g., "H", "Li", "Fe").
         electron_init_width: Width of Gaussian for electron initialization.
-        basis: Basis set name.
         ecp: Effective core potential name. Can be None (no ECP) or
             a string (e.g., "ccecp").
 
@@ -41,6 +39,5 @@ def atom_config(
         atoms=[atom],
         electron_spins=electron_spins,
         electron_init_width=electron_init_width,
-        basis=basis,
         ecp=ecp,
     )

src/jaqmc/app/molecule/config/base.py

@@ -2,12 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from dataclasses import field
+from typing import Literal
 
 from jaqmc.utils.atomic import Atom
 from jaqmc.utils.atomic.atomic_system import AtomicSystemConfig
+from jaqmc.utils.atomic.pretrain import PretrainReferenceConfig
 from jaqmc.utils.config import configurable_dataclass
 
-__all__ = ["MoleculeConfig"]
+__all__ = ["MoleculeConfig", "MoleculePretrainReferenceConfig"]
 
 
 @configurable_dataclass
@@ -16,3 +18,9 @@ class MoleculeConfig(AtomicSystemConfig):
     electron_spins: tuple[int, int] = (1, 0)
     fixed_spins_per_atom: list[tuple[int, int]] | None = None
     electron_init_width: float = 1.0
+
+
+@configurable_dataclass
+class MoleculePretrainReferenceConfig(PretrainReferenceConfig):
+    method: Literal["RHF", "UHF"] = "UHF"
+    "Variants of Hartree-Fock method."

src/jaqmc/app/molecule/config/diatomic.py

@@ -51,7 +51,6 @@ def diatomic_config(
     formula: str = "H2",
     bond_length: float = 1.4,
     unit: LengthUnit = LengthUnit.bohr,
-    basis: str | dict[str, str] = "sto-3g",
     ecp: str | dict[str, str] | None = None,
     spin: int = 0,
     electron_init_width: float = 1.0,
@@ -66,8 +65,6 @@ def diatomic_config(
         bond_length: Distance between the two atoms.
         unit: Length unit for ``bond_length`` and atom coordinates.
             Either ``"bohr"`` or ``"angstrom"``.
-        basis: Basis set name, or per-element mapping
-            (e.g., ``{"Li": "ccecpccpvdz", "H": "cc-pvdz"}``).
         ecp: Effective core potential specification. Can be ``None``
             (all-electron), a string (e.g., ``"ccecp"``), or a
             per-element mapping (e.g., ``{"Li": "ccecp"}``).
@@ -117,6 +114,5 @@ def diatomic_config(
         atoms=atoms,
         electron_spins=(n_alpha, n_beta),
         electron_init_width=electron_init_width,
-        basis=basis,
         ecp=ecp,
     )