def step_operator(self, f, lhs, rho):
from pystella import diff
D = diff(lhs, f)
from pymbolic import var
omega = var("omega")
return f - omega * (lhs - rho) / D
def step_operator(self, f, lhs, rho):
from pystella import diff
D = diff(lhs, f)
R_y = lhs - D * f # FIXME: only valid for linear equations
from pymbolic import var
omega = var("omega")
return (1 - omega) * f + omega * (rho - R_y) / D
def test_field_diff(proc_shape):
if proc_shape != (1, 1, 1):
pytest.skip("test field only on one rank")
from pystella import diff
y = ps.Field("y")
assert diff(y, y) == 1
assert diff(y[0], y[0]) == 1
assert diff(y[0], y[1]) == 0
y = ps.DynamicField("y")
assert diff(y, y) == 1
assert diff(y[0], y[0]) == 1
assert diff(y[0], y[1]) == 0
import pymbolic.primitives as pp
assert diff(y**3, y, "t") == pp.Product((3, 2, y, y.d(0)))
assert diff(y**3, "t", y) == pp.Product((3, y.d(0), 2, y))
for i, x in enumerate(["t", "x", "y", "z"]):
assert diff(y, x) == y.d(i)
assert diff(y[1, 3], x) == y.d(1, 3, i)
assert diff(y[1]**2, x) == 2 * y[1] * y.d(1, i)
dfdt[i] = df0[i]
# compute energy of background fields and initialize expansion
energy = compute_energy(f, dfdt, lap_f, dfdx, 1.)
expand = ps.Expansion(energy["total"], Stepper, mpl=mpl)
# compute hubble correction to scalar field effective mass
addot = expand.addot_friedmann_2(expand.a, energy["total"], energy["pressure"])
hubbleCorrection = -addot / expand.a
# effective masses of scalar fields
from pymbolic import var
from pymbolic.mapper.evaluator import evaluate_kw
fields = [var("f0")[i] for i in range(nscalars)]
d2Vd2f = [ps.diff(potential(fields), field, field) for field in fields]
eff_mass = [evaluate_kw(x, f0=f0) + hubbleCorrection for x in d2Vd2f]
modes = ps.RayleighGenerator(ctx,
fft,
dk,
volume,
seed=49279 * (decomp.rank + 1))
for fld in range(nscalars):
modes.init_WKB_fields(f[fld],
dfdt[fld],
norm=mphi**2,
omega_k=lambda k: np.sqrt(k**2 + eff_mass[fld]),
hubble=expand.hubble[0])