alienx — AlienX: a general fluid solver

The grid is dead. The manifold is awake. — from the Grimoire of Elbàlor

Rotation- and scale-invariant neural operators — AlienX 2D and AlienX 3D — that learn the time-step of any fluid from data. Architecture-only package: Darcy, Navier–Stokes, MHD, Gross–Pitaevskii, Rayleigh–Bénard, elasticity, or your own PDE. No checkpoints, no pretrained weights, no dataset. You train the weights on your own fields; every constructor and state_dict key is byte-compatible with the released paper code, so published checkpoints load into models built here.

pip install alienx

One front end, any fluid

AlienX is a general fluid solver: the architecture only ever sees a scalar field (whose gradient + Hessian build the local frame), two vector slots, and one positive regime number. alienx.fields is the outer shell that maps your raw fields onto those slots — Darcy, Navier–Stokes, MHD, Gross–Pitaevskii (2D or 3D), plasma, anything:

from alienx import (create_architecture, features_2d, features_2d_complex,
                    features_3d, features_3d_velocity)

# 2D scalar PDE (Darcy, porous media, reaction–diffusion, ...)
p = create_architecture("darcy2d")(**features_2d(k_field))

# 2D complex PDE (Gross–Pitaevskii, nonlinear Schrödinger, ...)
dr, di = create_architecture("gp2d")(**features_2d_complex(psi))

# 3D Navier–Stokes: paper channel construction (k = |curl u|, om, u, Re)
om = create_architecture("ns3d", re_max=2000)(**features_3d_velocity(u, re=900))

# 3D MHD: scalar = |B| builds the frame, slots = current J and velocity u,
# regime = magnetic Reynolds (any positive label)
out = create_architecture("ns3d", re_max=2000)(**features_3d(B.abs(), J, u, label=1e4))

# 3D Gross–Pitaevskii: scalar = |psi|^2, slots = probability current and flow
out = create_architecture("ns3d", re_max=2000)(**features_3d(k_psi, j, v, label=1.0))

Ground-truth generators for the shipped recipes live in alienx.solvers: Taylor–Green / IMEX Navier–Stokes (tgv_trajectory), 2D + 3D split-step Gross–Pitaevskii (gp_step, gp_step_3d, generate_gp_pair_3d), all with the same spectral convention as the feature builders. Helpers: curl_3d, biot_savart_3d (exact inverse curl: TGV roundtrip ~6e-7, div-free to 1e-15), probability_current, wavenumbers_2d/3d, spectral gradients/Hessians.

Every path is exercised by tests/smoke_fields.py on real training runs (2D GP residual fit, MHD induction step, 3D GP current step — all through the public model(**features(...)) API only).

The three operators

Quickstart

import torch
from alienx import (create_architecture, get_cached_iso_knn,
                    interior_mse, make_darcy_trainer, cosine_schedule)

model = create_architecture("darcy2d", hidden_dim=128, num_blocks=4)
trainer = make_darcy_trainer(model)            # paper recipe: AdamW 2e-3/1e-4
sched = cosine_schedule(500 * steps_per_epoch, lr=2e-3)

def loss_fn(model, K):
    coords, k, grad_k, grad_k_mag, p_true = next(darcy_batch)   # your data
    knn = get_cached_iso_knn(resolution, device)
    p_pred = model(coords, k, grad_k, grad_k_mag, knn)
    return interior_mse(p_pred, p_true, resolution)

trainer.fit(loss_fn, steps=len(sched), lr_schedule=sched, log_every=50)

3D uses the physics-informed loss (paper §5.3) and the K-unroll curriculum (K = 1→5 over the first 60 epochs):

from alienx import PhysicsLoss, make_ns_trainer

model = create_architecture("ns3d", re_max=2000.0)
trainer = make_ns_trainer(model)               # paper recipe: AdamW 3e-4, clip 1.0
physics = PhysicsLoss()                        # 1·MSE + 2·logE + 5·logZ + .05·div + .3·spec

def loss_fn(model, K):
    pred = model(coords, k, gk, gk_m, H, om, u, re_label, knn)  # your data
    u_pred = biot_savart(pred)                 # exact spectral recovery
    return physics(pred, omega_true, u_pred=u_pred, u_true=u_true,
                   KX=KX, KY=KY, KZ=KZ, K2=K2)

trainer.fit(loss_fn, steps=..., lr_schedule=cosine_schedule(...))

What the papers claim (and this code reproduces)

• 2D: interior L1 0.0036–0.0087 across 16²→256² zero-shot; rotation sweep flat (arbitrary angles within 0.003 of baseline); bitwise-exact D4 (pred(R·k) = R·pred(k)); 45° residual eliminated by the isotropic 24-neighbor stencil.
• 2D × GP (cross-PDE): 0.34% single-step rel RMSE on Gross–Pitaevskii with only outer-shell substitutions; phase ablation 108.6×.
• 3D: 0.46% single-step rel RMSE across Re ∈ {300, 628, 1200}; zero-shot 0.53% at held-out Re = 900; QSA ablation 32.81×; phase-only ablation 4.74×; 73-step rollout, energy drift < 6%; divergence-free by construction (~2.7e−6 float32 round-off).

Requirements

Python ≥ 3.9, torch ≥ 2.0, numpy. CPU is enough for 2D; 3D training wants a GPU (paper timings: T4, ~18 h for the full K-curriculum at G=32).

References

1. Yaka, E. H. D. AlienX (ISN): A Continuous Rotation- and Scale-Invariant Operator on Geometric Manifolds.
2. Yaka, E. H. D. AlienX 3D: A General SO(3)-Equivariant Neural Operator for 3D Physical Fields.

Model cards + papers: ElBalor/AlienX-2D-Isotropic-Stencil · ElBalor/AlienX-3D-ISN-Operator

License

AGPL-3.0.

Eric Yaka (Elbàlor / The Digital Necromancer) · Capital Software · Abuja, Nigeria