Reference

wavefunction module

class wavefunction.BosonicWavefunction(*args, **kwargs)

Bases: Module

A wavefunction of many indistinguishable bosons.

coords: SumCoordinates
mlp: MLP
class wavefunction.ConfinedWavefunction(*args, **kwargs)

Bases: Module

A wavefunction naturally confined in a harmonic potential.

Parameters:

  • wf –

  • L –

Returns:

A wavefunction

L: float
wf: Module
class wavefunction.DetCoordinate(*args, **kwargs)

Bases: Module

Totally antisymmetric coordinates, constructed by taking a determinant.

A single determinental coordinate, of N particles, is defined by constructing N one-particle coordinates, and creating a matrix in which element (i,j) contains the i`th coordinate of particle `j. The determinant of that matrix is totally antisymmetric under particle exchange.

D: int
N: int
fs: list
class wavefunction.DetCoordinates(*args, **kwargs)

Bases: Module

D: int
K: int
N: int
coords: list
class wavefunction.DistinguishableWavefunction(*args, **kwargs)

Bases: Module

A distinguishable many-body wavefunction.

Underneath the hood, this is not really different from OneParticleWavefunction, except that convenience reshaping is performed.

D: int
N: int
mlp: MLP
class wavefunction.FermionicWavefunction(*args, **kwargs)

Bases: Module

A wavefunction of many indistinguishable fermions.

coords: DetCoordinates
mlp: OddMLP
class wavefunction.GaussianWavefunction(*args, **kwargs)

Bases: Module

A simple Gaussian. Mostly for testing purposes.

D: int
N: int
class wavefunction.MLP(*args, **kwargs)

Bases: Module

static activation(x: Any, alpha: Any = 1.0) → Any

Continuously-differentiable exponential linear unit activation.

Computes the element-wise function:

\[\begin{split}\mathrm{celu}(x) = \begin{cases} x, & x > 0\\ \alpha \left(\exp(\frac{x}{\alpha}) - 1\right), & x \le 0 \end{cases}\end{split}\]

For more information, see Continuously Differentiable Exponential Linear Units.

Parameters:

  • x – input array

  • alpha – array or scalar (default: 1.0)

layers: list
class wavefunction.ManyBodyWavefunction(*args, **kwargs)

Bases: Module

A wavefunction of many particles, with arbitrary distinguishability.

Each species is labeled by a string.

class wavefunction.OddMLP(*args, **kwargs)

Bases: Module

A multi-layer perceptron with the guarantee that f(x) = -f(-x).

static activation(x)

Compute hyperbolic tangent element-wise.

LAX-backend implementation of numpy.tanh().

Original docstring below.

Equivalent to np.sinh(x)/np.cosh(x) or -1j * np.tan(1j*x).

Parameters:

x (array_like) – Input array.

Returns:

y – The corresponding hyperbolic tangent values. This is a scalar if x is a scalar.

Return type:

ndarray

References

layers: list
class wavefunction.OneParticleWavefunction(*args, **kwargs)

Bases: Module

Wavefunction of a single particle, in D dimensions.

mlp: MLP
class wavefunction.RadialWavefunction(*args, **kwargs)

Bases: Module

A many-body wavefunction that depends only on the radial coordinates and distances.

D: int
N: int
mlp: MLP
class wavefunction.RelativeCoordinates(*args, **kwargs)

Bases: Module

Only suitable for distinguishable particles.

class wavefunction.SumCoordinate(*args, **kwargs)

Bases: Module

D: int
N: int
f: MLP
class wavefunction.SumCoordinates(*args, **kwargs)

Bases: Module

D: int
K: int
N: int
coords: list
class model.ElectricModel(alpha, masses, charges, species, confining=1.0)

Bases: ManyBodyModel

potential(x)
Parameters:

x – A tensor of shape (N,D).

class model.FermiGasModel(D, N, L, a, g)

Bases: ManyBodyModel

potential(x)
Parameters:

x – A tensor of shape (N,D).

sampler(key, wf)
Parameters:

  • key –

  • wf –

  • K – batch size

Returns:

A function that yields samples.

wavefunction(key)

Initialize and return a wavefunction.

Parameters:

key –

Returns:

class model.ManyBodyModel

Bases: Model

hamiltonian(wf, x)
kinetic(wf, x)
abstract potential(x)
Parameters:

x – A tensor of shape (N,D).

class model.Model

Bases: ABC

abstract sampler(key, wf, K)
Parameters:

  • key –

  • wf –

  • K – batch size

Returns:

A function that yields samples.

abstract wavefunction()

Initialize and return a wavefunction.

Parameters:

key –

Returns:

class model.OscillatorModel(D, c2, c4)

Bases: Model

hamiltonian(wf, x)
particles = None
potential(x)
sampler(key, wf, N)
Parameters:

  • key –

  • wf –

  • K – batch size

Returns:

A function that yields samples.

wavefunction(key)

Initialize and return a wavefunction.

Parameters:

key –

Returns:

class model.SigmaOmegaModel(particles, confining=20.0)

Bases: ManyBodyModel

potential(x)
Parameters:

x – A tensor of shape (N,D).

sampler(key, wf)
Parameters:

  • key –

  • wf –

  • K – batch size

Returns:

A function that yields samples.

wavefunction(key)

Initialize and return a wavefunction.

Parameters:

key –

Returns: