1import numpy as np
 2from neural_levy_noise import compensated_euler_step, small_variance, small_mean
 3
 4rng=np.random.default_rng(7)
 5x=np.zeros((4096,3)); d=np.ones_like(x)
 6y,k=compensated_euler_step(x,d,0.1,0.05,1.5,rng=rng)
 7assert y.shape==x.shape and isinstance(k,int) and k>0
 8# Repeated scalar noise checks matched per-coordinate Gaussian variance after removing drift/jumps
 9assert abs(small_variance(.05,1.5)-0.5**0) > 0 # formula callable
10assert small_mean(.05,.5)>0
11try:
12    small_mean(.05,1.0)
13    raise AssertionError('expected divergent mean guard')
14except ValueError:
15    pass
16# one-sided corrected update also runs
17z,k2=compensated_euler_step(np.zeros((16,2)),np.zeros((16,2)),.01,.1,.5,rng=rng,one_sided=True,add_small_mean=True)
18assert z.shape==(16,2) and k2>=0
19print('smoke_ok', y.shape, k, k2, 'variance', small_variance(.05,1.5))