2D Turbulence

import mmf_setup;mmf_setup.nbinit()
from gpe.imports import *

This cell adds /home/docs/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src to your path, and contains some definitions for equations and some CSS for styling the notebook. If things look a bit strange, please try the following:

  • Choose "Trust Notebook" from the "File" menu.
  • Re-execute this cell.
  • Reload the notebook.

[I 19:59:31 numexpr.utils] NumExpr defaulting to 2 threads.
---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[1], line 2
      1 import mmf_setup;mmf_setup.nbinit()
----> 2 from gpe.imports import *

File ~/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src/gpe/imports.py:28
     26 from mmfutils.plot import imcontourf  # noqa: E402
     27 from gpe.minimize import MinimizeState  # noqa: E402
---> 28 from gpe.utils import evolve_to, evolve, evolves  # noqa: E402
     29 from gpe.plot_utils import MPLGrid  # noqa: E402
     30 from pytimeode.evolvers import EvolverSplit, EvolverABM  # noqa: E402

File ~/checkouts/readthedocs.org/user_builds/gpe/checkouts/latest/src/gpe/utils.py:35
     32 from persist.objects import Archivable
     33 from persist.archive import Archive
---> 35 from pytimeode.evolvers import EvolverABM
     36 from pytimeode.mixins import ArrayStateMixin
     37 from pytimeode.interfaces import implementer, IStateForABMEvolvers

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/pytimeode/__init__.py:2
      1 from . import interfaces
----> 2 from . import mixins
      3 from . import evolvers
      5 __all__ = ["interfaces", "mixins", "evolvers"]

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/pytimeode/mixins.py:39
     30 from .interfaces import (
     31     IState,
     32     IStateApply,
   (...)     36     implementer,
     37 )
     38 from . import interfaces
---> 39 from .utils import expr
     41 __all__ = [
     42     "StateMixin",
     43     "StatesMixin",
   (...)     47     "ArraysStateWithBraketMixin",
     48 ]
     51 @implementer(IState)
     52 class StateMixin:

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/pytimeode/utils/expr.py:18
     15 import numpy as np
     17 try:
---> 18     import sympy
     19 except ImportError:
     20     sympy = None

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/__init__.py:77
     70 from .logic import (to_cnf, to_dnf, to_nnf, And, Or, Not, Xor, Nand, Nor,
     71         Implies, Equivalent, ITE, POSform, SOPform, simplify_logic, bool_map,
     72         true, false, satisfiable)
     74 from .assumptions import (AppliedPredicate, Predicate, AssumptionsContext,
     75         assuming, Q, ask, register_handler, remove_handler, refine)
---> 77 from .polys import (Poly, PurePoly, poly_from_expr, parallel_poly_from_expr,
     78         degree, total_degree, degree_list, LC, LM, LT, pdiv, prem, pquo,
     79         pexquo, div, rem, quo, exquo, half_gcdex, gcdex, invert,
     80         subresultants, resultant, discriminant, cofactors, gcd_list, gcd,
     81         lcm_list, lcm, terms_gcd, trunc, monic, content, primitive, compose,
     82         decompose, sturm, gff_list, gff, sqf_norm, sqf_part, sqf_list, sqf,
     83         factor_list, factor, intervals, refine_root, count_roots, all_roots,
     84         real_roots, nroots, ground_roots, nth_power_roots_poly, cancel,
     85         reduced, groebner, is_zero_dimensional, GroebnerBasis, poly,
     86         symmetrize, horner, interpolate, rational_interpolate, viete, together,
     87         BasePolynomialError, ExactQuotientFailed, PolynomialDivisionFailed,
     88         OperationNotSupported, HeuristicGCDFailed, HomomorphismFailed,
     89         IsomorphismFailed, ExtraneousFactors, EvaluationFailed,
     90         RefinementFailed, CoercionFailed, NotInvertible, NotReversible,
     91         NotAlgebraic, DomainError, PolynomialError, UnificationFailed,
     92         GeneratorsError, GeneratorsNeeded, ComputationFailed,
     93         UnivariatePolynomialError, MultivariatePolynomialError,
     94         PolificationFailed, OptionError, FlagError, minpoly,
     95         minimal_polynomial, primitive_element, field_isomorphism,
     96         to_number_field, isolate, round_two, prime_decomp, prime_valuation,
     97         galois_group, itermonomials, Monomial, lex, grlex,
     98         grevlex, ilex, igrlex, igrevlex, CRootOf, rootof, RootOf,
     99         ComplexRootOf, RootSum, roots, Domain, FiniteField, IntegerRing,
    100         RationalField, RealField, ComplexField, PythonFiniteField,
    101         GMPYFiniteField, PythonIntegerRing, GMPYIntegerRing, PythonRational,
    102         GMPYRationalField, AlgebraicField, PolynomialRing, FractionField,
    103         ExpressionDomain, FF_python, FF_gmpy, ZZ_python, ZZ_gmpy, QQ_python,
    104         QQ_gmpy, GF, FF, ZZ, QQ, ZZ_I, QQ_I, RR, CC, EX, EXRAW,
    105         construct_domain, swinnerton_dyer_poly, cyclotomic_poly,
    106         symmetric_poly, random_poly, interpolating_poly, jacobi_poly,
    107         chebyshevt_poly, chebyshevu_poly, hermite_poly, hermite_prob_poly,
    108         legendre_poly, laguerre_poly, apart, apart_list, assemble_partfrac_list,
    109         Options, ring, xring, vring, sring, field, xfield, vfield, sfield)
    111 from .series import (Order, O, limit, Limit, gruntz, series, approximants,
    112         residue, EmptySequence, SeqPer, SeqFormula, sequence, SeqAdd, SeqMul,
    113         fourier_series, fps, difference_delta, limit_seq)
    115 from .functions import (factorial, factorial2, rf, ff, binomial,
    116         RisingFactorial, FallingFactorial, subfactorial, carmichael,
    117         fibonacci, lucas, motzkin, tribonacci, harmonic, bernoulli, bell, euler,
   (...)    138         Znm, elliptic_k, elliptic_f, elliptic_e, elliptic_pi, beta, mathieus,
    139         mathieuc, mathieusprime, mathieucprime, riemann_xi, betainc, betainc_regularized)

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/polys/__init__.py:79
      3 __all__ = [
      4     'Poly', 'PurePoly', 'poly_from_expr', 'parallel_poly_from_expr', 'degree',
      5     'total_degree', 'degree_list', 'LC', 'LM', 'LT', 'pdiv', 'prem', 'pquo',
   (...)     65     'field', 'xfield', 'vfield', 'sfield'
     66 ]
     68 from .polytools import (Poly, PurePoly, poly_from_expr,
     69         parallel_poly_from_expr, degree, total_degree, degree_list, LC, LM,
     70         LT, pdiv, prem, pquo, pexquo, div, rem, quo, exquo, half_gcdex, gcdex,
   (...)     76         nth_power_roots_poly, cancel, reduced, groebner, is_zero_dimensional,
     77         GroebnerBasis, poly)
---> 79 from .polyfuncs import (symmetrize, horner, interpolate,
     80         rational_interpolate, viete)
     82 from .rationaltools import together
     84 from .polyerrors import (BasePolynomialError, ExactQuotientFailed,
     85         PolynomialDivisionFailed, OperationNotSupported, HeuristicGCDFailed,
     86         HomomorphismFailed, IsomorphismFailed, ExtraneousFactors,
   (...)     91         MultivariatePolynomialError, PolificationFailed, OptionError,
     92         FlagError)

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/polys/polyfuncs.py:10
      8 from sympy.polys.polyoptions import allowed_flags, build_options
      9 from sympy.polys.polytools import poly_from_expr, Poly
---> 10 from sympy.polys.specialpolys import (
     11     symmetric_poly, interpolating_poly)
     12 from sympy.polys.rings import sring
     13 from sympy.utilities import numbered_symbols, take, public

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/polys/specialpolys.py:298
    294     return dmp_mul(f, h, n, K), dmp_mul(g, h, n, K), h
    296 # A few useful polynomials from Wang's paper ('78).
--> 298 from sympy.polys.rings import ring
    300 def _f_0():
    301     R, x, y, z = ring("x,y,z", ZZ)

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/polys/rings.py:31
     27 from sympy.polys.polyoptions import (Domain as DomainOpt,
     28                                      Order as OrderOpt, build_options)
     29 from sympy.polys.polyutils import (expr_from_dict, _dict_reorder,
     30                                    _parallel_dict_from_expr)
---> 31 from sympy.printing.defaults import DefaultPrinting
     32 from sympy.utilities import public, subsets
     33 from sympy.utilities.iterables import is_sequence

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/printing/__init__.py:11
      7 from .mathml import mathml, print_mathml
      9 from .python import python, print_python
---> 11 from .pycode import pycode
     13 from .codeprinter import print_ccode, print_fcode
     15 from .codeprinter import ccode, fcode, cxxcode, rust_code # noqa:F811

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/printing/pycode.py:11
      9 from sympy.core.mod import Mod
     10 from .precedence import precedence
---> 11 from .codeprinter import CodePrinter
     13 _kw = {
     14     'and', 'as', 'assert', 'break', 'class', 'continue', 'def', 'del', 'elif',
     15     'else', 'except', 'finally', 'for', 'from', 'global', 'if', 'import', 'in',
     16     'is', 'lambda', 'not', 'or', 'pass', 'raise', 'return', 'try', 'while',
     17     'with', 'yield', 'None', 'False', 'nonlocal', 'True'
     18 }
     20 _known_functions = {
     21     'Abs': 'abs',
     22     'Min': 'min',
     23     'Max': 'max',
     24 }

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/printing/codeprinter.py:13
     11 from sympy.core.sorting import default_sort_key
     12 from sympy.core.symbol import Symbol
---> 13 from sympy.functions.elementary.complexes import re
     14 from sympy.printing.str import StrPrinter
     15 from sympy.printing.precedence import precedence, PRECEDENCE

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/functions/__init__.py:28
     25 from sympy.functions.elementary.integers import floor, ceiling, frac
     26 from sympy.functions.elementary.piecewise import (Piecewise, piecewise_fold,
     27                                                   piecewise_exclusive)
---> 28 from sympy.functions.special.error_functions import (erf, erfc, erfi, erf2,
     29         erfinv, erfcinv, erf2inv, Ei, expint, E1, li, Li, Si, Ci, Shi, Chi,
     30         fresnels, fresnelc)
     31 from sympy.functions.special.gamma_functions import (gamma, lowergamma,
     32         uppergamma, polygamma, loggamma, digamma, trigamma, multigamma)
     33 from sympy.functions.special.zeta_functions import (dirichlet_eta, zeta,
     34         lerchphi, polylog, stieltjes, riemann_xi)

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/functions/special/error_functions.py:22
     20 from sympy.functions.elementary.hyperbolic import cosh, sinh
     21 from sympy.functions.elementary.trigonometric import cos, sin, sinc
---> 22 from sympy.functions.special.hyper import hyper, meijerg
     24 # TODO series expansions
     25 # TODO see the "Note:" in Ei
     26 
     27 # Helper function
     28 def real_to_real_as_real_imag(self, deep=True, **hints):

File ~/checkouts/readthedocs.org/user_builds/gpe/conda/latest/lib/python3.14/site-packages/sympy/functions/special/hyper.py:978
    973     @classmethod
    974     def _expr_big_minus(cls, z, n):
    975         return S.NegativeOne**n*(asinh(sqrt(z))/sqrt(z) + n*pi*I/sqrt(z))
--> 978 class HyperRep_asin2(HyperRep):
    979     """ Represent hyper([1, 1], [3/2], z) == asin(sqrt(z))/sqrt(z)/sqrt(1-z). """
    980     # TODO this can be nicer

KeyboardInterrupt: 

2D Turbulence#

Here we reproduce the results of [Zhao et al., 2025], where they stir a 2D BEC and claim to measure a Kolmogorov decay spectrum.

  • \(N = 2\times 10^{5}\) atoms.

  • \(R = r_{\mathrm{tr}} = \qty{22}{micron}\)

from gpe.imports import *
from gpe.bec import StateBase, u
from gpe.utils import ExperimentBase, get_good_N, _GPU
from gpe.minimize import MinimizeState

from mmfutils.math.special import mstep

@_GPU.add_non_GPU_methods
class State(StateBase):
    def __init__(self, experiment, kcut, kappa0, **kw):
        self.experiment = experiment
        self.kcut = kcut
        self.kappa0 = kappa0
        super().__init__(**kw)

    def init(self):
        super().init()

        ks = s.basis.kx
        mask = np.logical_and(ks <= self.kcut, ks >= -self.kcut)
        self._phase_k = np.exp(-1j * self.kappa0 * np.where(mask, 0, 1))
        
    def get_Vext_GPU(self):
        V_ext = self.experiment.get_Vext(state=self)
        if (self.initializing or self.t < 0) and getattr(self, "mu", None):
            V_ext -= self.mu
        return V_ext

    def _apply_dissipation(self, dy):
        """Return applying dissipation."""
        fftn, ifftn = self.basis.fftn, self.basis.ifftn
        dy.set_psi(ifftn(self._phase_k * fftn(dy.get_psi())))
        
    def compute_dy_dt(self, dy, subtract_mu=True):
        super().compute_dy_dt(dy, subtract_mu=subtract_mu)
        if not (self.initializing or self.t < 0):
            self._apply_dissipation(dy)
        return dy


class Experiment(ExperimentBase):
    hbar = 1
    m = u.m_Rb87
    w_z = 2*np.pi * (220*u.Hz)
    R_micron = 22.0
    R_Rx = 0.9
    g = 1
    V0_mu = 100.0
    V_dr_healing_length = 1.0

    mu = None
    healing_length_micron = 2.0
    dx_healing_length = 0.1
    n0_micron2 = 200.0   # Background density
    
    # Stirring Potential
    Vstir_Hz = 25
    Vstir_R_micron = 3.5
    Vstir_dr_healing_length = 1.0
    Vstir_mu = 10.0
    t_stir_ms = 16
    rng_seed = 0

    # Dissipation
    kcut_invmicron = 5
    kappa0 = 0.02

    t_final_ms = 56
    
    State = State
    
    def init(self):
        if self.mu is not None and self.healing_length_micron is None:
            self.healing_length = np.sqrt(self.hbar**2 / 2 / self.m / self.mu)
            mu = self.mu

        elif self.mu is None and self.healing_length_micron is not None:
            self.healing_length = self.healing_length_micron * u.micron
            mu = self.hbar**2 / 2 / self.m / self.healing_length**2

        else:
            raise ValueError("Must specify (only) one of `mu` or `healing_length_micron`")

        dx = self.dx_healing_length * self.healing_length
        self.V_R = self.R_micron * u.micron
        Rx = self.V_R / self.R_Rx
        Nx = get_good_N(2*Rx/dx)
        n0 = self.n0_micron2 / u.micron**2

        kcut = self.kcut_invmicron / u.micron
        
        g = mu / n0
        self.state_args = dict(kcut=kcut, kappa0=self.kappa0, Nxyz=(Nx, Nx), Lxyz=(2*Rx, 2*Rx),
                               g=g, mu=mu, m=self.m, hbar=self.hbar)
        self.V0 = self.V0_mu * mu
        self.V_dr = self.V_dr_healing_length * self.healing_length
        self._Vtrap = None
        
        self.Vstir_w = 2*np.pi * self.Vstir_Hz * u.Hz
        self.Vstir_R = self.Vstir_R_micron * u.micron
        self.Vstir_dr = self.Vstir_dr_healing_length * self.healing_length
        self.Vstir_V0 = self.Vstir_mu * mu

        self._n_R_changes = 0
        self.rng = np.random.default_rng(seed=self.rng_seed)
        self._R_stir = 13.5 * u.micron  # initial radius of stirring

        self.t_stir = self.t_stir_ms * u.ms
        self.t_final = self.t_final_ms * u.ms

    def get_Vext(self, state):
        """Return the external potential."""
        if self._Vtrap is None:
            self._Vtrap = self.get_Vtrap(state=state)
        V_ext = self._Vtrap + self.get_Vstir(state=state)
        return V_ext
        
    def get_Vtrap(self, state):
        """Return the static trapping potential."""
        x, y = state.get_xyz()
        r = np.sqrt(x**2 + y**2)
        return self._Vstep(r, R=self.V_R, dr=self.V_dr)
        
    def _Vstep(self, r, R, dr):
        """Smooth step at R of width dr."""
        return mstep(r - (R - dr/2), dr)

    def _Vrod(self, r):
        """Stirring rod potential (central)."""
        return 1-self._Vstep(r, R=self.Vstir_R, dr=self.Vstir_dr)

    def get_Vstir(self, state):
        """Return the dynamic stirring potential."""
        t = state.t
        if t > self.t_stir:
            return 0

        th0 = -np.pi/4
        dth = self.Vstir_w * t

        # TODO: Need to make these random transitions smooth?
        if t >= 400 * u.s * 1e-6 * self._n_R_changes:
            self._n_R_changes += 1
            self._R_stir = self.rng.uniform(12, 15) * u.micron
        
        z0, z1 = self._R_stir * np.exp([1j*(th0 + dth), 1j*(th0 - dth)])
        x, y = state.get_xyz()
        z = x + 1j*y
        return self.Vstir_V0 * (self._Vrod(abs(z-z0)) + self._Vrod(abs(z-z1)))
        
    def get_state(self):
        state = State(experiment=self, **self.state_args)
        return state
        
    def get_initialized_state(self):
        state0 = self.get_state()
        m = MinimizeState(state0)
        print(m.check())
        state1 = m.minimize(use_scipy=True)
        state = self.get_state()
        state.set_psi(state1.get_psi())
        return state

e = Experiment()
s = e.get_initialized_state()
s.plot()
from pytimeode.evolvers import EvolverABM
from mmfutils.contexts import FPS

ev = EvolverABM(s, dt=0.1*s.t_scale)

steps = 200
loops = int(e.t_final / ev.dt / steps)

for frame in FPS(loops, timeout=10):
    ev.evolve(steps)
    plt.clf()
    ev.y.plot()
    clear_output(wait=True)
    display(plt.gcf())
s0 = e.get_initialized_state()
rng = np.random.default_rng(seed=2)
def get_psi(s0=s0, p=0.1, steps=200, rng=rng):
    s = e.get_state()
    s.set_psi(s0.get_psi() * np.exp(1j*p*rng.normal(size=s0.shape)))
    ev = EvolverABM(s, dt=0.2*s.t_scale)
    ev.evolve(steps)
    return ev.y.get_psi()

psis = [get_psi() for n in range(10)]
n_ = np.abs(np.mean(psis, axis=0))**2
print(s0.get_N(), s0.integrate(n_))
#n = get_n()
#plt.imshow(n)