RFC: Ergonomic API for Continuous-Space Hilbert Spaces

[gcarleo]

Continuous-Space Hilbert API — design draft

This note proposes a small, self-contained layer for constructing NetKet Hilbert spaces that represent particles moving in continuous space. It is meant to cover molecules, electron gases, helium droplets, crystals, and other many-body systems without bringing in external quantum-chemistry packages. No Hamiltonians or operators are discussed here—only the Hilbert side.

---

1. Core data objects

Name	What it stores	Why it exists
ParticleSpecies	mass, charge, total spin S, optional label	immutable physical constants
Particle	reference to a ParticleSpecies; optional position position: jax.Array | None (bohr). None ⇒ coordinate is dynamical, otherwise fixed (Born–Oppenheimer). Optional spin projection m_z (None ⇒ spin dynamical).	one body
ParticleSet (sub-class of nk.hilbert.TensorHilbert)	takes Sequence[Particle] and optional geometry: nk.geometry.Cell | None; builds the tensor-product Hilbert in its constructor, so the instance is itself a Hilbert object. Also exposes stacked JAX arrays for positions, charges, masses, spin masks, etc.	single object that is both data container and Hilbert space

---

Arrays are stored in atomic units (bohr, electron-mass, elementary-charge); no unit objects appear in JIT-traced code.

2. Predefined species

Two constants cover the majority of use-cases:

Electron = ParticleSpecies(
    mass   = 1.0,              # m_e
    charge = -1.0,             # −e
    S      = 0.5,
    label  = "e⁻",
)

Proton = ParticleSpecies(
    mass   = 1836.15267389,    # CODATA-18
    charge = +1.0,
    S      = 0.5,
    label  = "p⁺",
)

Users can always define their own, e.g.

He4 = ParticleSpecies(mass=7294.30, charge=0.0, S=0.0, label="⁴He").

3. Turning particles into a Hilbert space

class ParticleSet(nk.hilbert.TensorHilbert):
    def __init__(self,
                 particles: Sequence[Particle],
                 geometry: nk.geometry.Cell | None = None):

        self.particles = list(particles)
        self.geometry  = geometry or nk.geometry.Cell(d=3, L=np.inf, pbc=False)

        blocks: list[nk.hilbert.AbstractHilbert] = []
        d = self.geometry.d

        for p in self.particles:
            # ── coordinate DOFs (always present) ───────────────────────
            if p.position is None:
                # dynamical coordinate: one ℝ site per spatial dimension
                blocks.extend([ContinuousSpace()] * d)
            else:
                # fixed coordinate: 1-dimensional delta Hilbert
                blocks.append(FixedPosition(value=p.position, dim=d))

            # ── spin DOF (always tensor-producted) ─────────────────────
            if p.species.S > 0.0:
                if p.m_z is None:
                    blocks.append(nk.hilbert.Spin(s=p.species.S))
                else:
                    blocks.append(FixedSpin(value=p.m_z))

        super().__init__(blocks)

ContinuousSpace — unbounded ℝ site.
FixedPosition — trivial continuos-space Hilbert peaked at a fixed coordinate (to be added).
FixedSpin — trivial 1-dimensional Hilbert holding a constant spin eigenvalue (to be added).
By delegating spin blocks to nk.hilbert.Spin we reuse its random-state and enumeration machinery unchanged.

4. Helper builders (optional sugar)

Each helper returns a ParticleSet, they serve as a convenient interface for commonly used cases, such as creating the Hilbert space for molecules or solids.

atoms = "H 0 0 0; H 0 0 1.4"
H2 = Molecule.from_string(atoms, units="bohr", charge=0, total_spin=0, quantum_nuclei=False) 
# parses "H 0 0 0; H 0 0 1.4" and adds electrons.

Crystal.from_unit_cell(atoms, cell_vectors, electrons_per_cell, units) 
# single-cell ParticleSet; translational symmetry handled elsewhere.

These wrappers are thin; they do not hide the underlying Particle/ParticleSet objects.

5. Units

API accepts "bohr" (alias "au") and "angstrom" (alias "Å").

Positions are converted once to bohr; internal arrays stay unit-free.

ParticleSet.to_units("angstrom") is available for inspection/export.

6. Usage examples

# 1 — Electron gas, 14 e⁻, periodic cube 10 bohr
geo = nk.geometry.Cell(d=3, L=10.0, pbc=True)
hilb = ParticleSet([Particle(Electron) for _ in range(14)], geometry=geo)

# 2 — Fixed-spin electron gas (7 ↑, 7 ↓)
up   = [Particle(Electron, m_z=+1) for _ in range(7)]
down = [Particle(Electron, m_z=-1) for _ in range(7)]
hilb_fs = ParticleSet(up + down, geometry=geo)

# 3 — H₂ molecule in free space (helpers attach geometry=None → free)
hilb_H2 = nkx.Molecule.from_string(
              "H 0 0 -0.7; H 0 0 0.7", units="bohr")

# 4 — Diamond Si unit cell
cell_vec = [[0,2.715,2.715],
            [2.715,0,2.715],
            [2.715,2.715,0]]  # Å
hilb_Si = nkx.Crystal.from_unit_cell(
              atoms=..., cell_vectors=cell_vec,
              electrons_per_cell=8, units="angstrom")

# 5 — Liquid ⁴He cluster, 32 spinless bosons
hilb_He = ParticleSet([Particle(He4) for _ in range(32)])

7. Rationale

Small surface – three fundamental data classes plus syntax sugar.

Spin always tensor-producted — fixed projections reduce to 1-dim blocks.

Reuse, not rewrite – spin blocks rely on nk.hilbert.Spin.

JAX-friendly – positions & masks are plain JAX arrays

Model-agnostic – mechanism handles molecules, fluids, crystals, etc.;
helpers are convenience only.

8. Open questions

Default initial electron positions: always origin, or randomised near nuclei?

Is a built-in isotope table worth the complexity, or should users specify masses when needed?

[gcarleo]

@PhilipVinc @dalin27 @gpescia @Z-Denis @inailuig @jwnys and any other who has worked / has interest in the continuous space interface, provide feedback on this proposal.

[jwnys]

This looks quite similar to my own code:

• I also have geometry as a separate class
• It would indeed be useful to distinguish the different species
• It's not clear to me here how you would fix the nuclei in molecules? So Particle should have a way to determine whether e.g. the position is fixed or not?
• In my opinion "quantum" flag is rather part of an operator and model, not so much the hilbert space object?

[gcarleo]

yes separating geometry and having multispeicies is among the fundamental requirements... we had discussions in some issue here two years ago already but never actually implemented it peoperly. now we can't really postpone much more I think...

in any case, the quantum keyword is essentially what tells if positions are static or not. could be improved, I am not 100 percent sure of that either

[gpescia]

Looks good.
Maybe some comments:

• We should make sure that 1+2 d is also supported/inferred from the positions.
• For molecules I saw improvements when initializing the right number of electrons for a given nucleus at position R, for example Z=6, 6 electrons around R. (as you say)
• For crystals there should be a possibility to say which atom is at which lattice site, I don't see how that would be done here or is that nkx.Crystal.from_unit_cell(...).particles ?
• Is there a distinction between primitive cell and super/simulation cell in geometry? For example simulating 1x1x1 (one primitive cell) vs 2x2x2 (2 primitive cells in each direction).

[gcarleo]

OK so a solution would be to set position=None, similarly to the m_z=None case, meaning in both cases that position and spin (S_z) are sampled. Instead if one specifies position=(0,0,2) or something, this is not going to be sampled. We can then remove the quantum keywork. I am going to update this.

[gcarleo]

For crystals there should be a possibility to say which atom is at which lattice site, I don't see how that would be done here or is that nkx.Crystal.from_unit_cell(...).particles ?

@gpescia yes indeed, I didn't discuss it but similarly to the helper to construct molecules there will be another one that will assign atoms to specific lattice sites

[gcarleo]

notice though that the molecule constructor and similar would fix those positions in the BO case, not Particle itself, so in this sense maybe we can just remove quantum as a keyword... then the Hamiltonian will work with those objects and know if it's BO or not

[gcarleo]

I removed the quantum keyword and simplified things...