This demo implements algorithms for stateless, constant-memory
sampling of entire functions from a Brownian motion of
any dimension. This kind of sampler is useful in adaptive
stochastic partial differential equation solvers, since it
allows the noise to be sampled in a stateless manner.
This code also demonstrates Dex's ability to do fast stateful for loops,
split random keys in a fine-grained way,
and how to use the typeclass system to implement
a form of recursion.
One-dimensional Brownian Motion on the Unit Interval
The rest of the algorithms in this file simply translate and
scale calls to this function.
structSamplerState(v)=k:Keyy:vsigma:Floatt:Float
defsample_unit_brownian_bridge(tolerance:Float,sampler:(Key)->v,key:Key,t:Float)->vgiven(v|VSpace)=-- Can only handle t between 0 and 1.
-- iteratively subdivide to desired tolerance.num_iters=10+f_to_n(-logtolerance)init_state=SamplerState(key,zero,1.0,t)result=foldinit_state\i:(Finnum_iters)prev.[key_draw,key_left,key_right]=split_keykey-- add scaled noiset'=abs(prev.t-0.5)new_y=prev.y+prev.sigma*(0.5-t').*samplerkey_draw-- zoom in left or rightnew_key=ifprev.t>0.5thenkey_leftelsekey_rightnew_sigma=prev.sigma/sqrt2.0new_t=t'*2.0SamplerState(new_key,new_y,new_sigma,new_t)result.y
One-dimensional Brownian motion anywhere on the real line
These functions generalize the sampling of a Brownian bridge on the unit
interval to sample from a Brownian motion defined anywhere on the real line.
When evaluating at times larger than 1, an iterative doubling procedure
eventually produces values bracketing the desired time.
These bracketing values are then used to call the iterative
Brownian bridge sampler.
defsample_brownian_bridge(tolerance:Float,sampler:(Key)->v,key:Key,t0:Float,y0:v,t1:Float,y1:v,t:Float)->vgiven(v|VSpace)=-- Linearly interpolate between the two bracketing samples
-- and add appropriately scaled Brownian bridge noise.unit_bb=\t.sample_unit_brownian_bridge(tolerance/(t1-t0))samplerkeytbridge_val=scale_brownian_bridgeunit_bbt0t1tinterp_val=linear_interpt0y0t1y1tbridge_val+interp_val
defsample_brownian_motion(tolerance:Float,sampler:(Key)->v,key:Key,t':Float)->vgiven(v|VSpace)=-- Can handle a t' anywhere on the real line.
-- Handle negative times by reflecting and using a different key.[neg_key',key']=split_keykey(key,t)=ift'<0.0then(neg_key',-t')else(key',t')[key_start,key_rest]=split_keykeyfirst_f=samplerkey_start-- Iteratively sample a point twice as far ahead
-- until we bracket the query time.doublings=Fin(1+(natlog2(f_to_nt)))init_state=(0.0,zero,1.0,first_f,key_rest)(t0,left_f,t1,right_f,key_draw)=foldinit_state\i:doublingsstate.(t0,left_f,t1,right_f,key_inner)=state[key_current,key_continue]=split_keykey_restnew_right_f=left_f+(sqrtt1).*samplerkey_current(t1,right_f,t1*2.0,new_right_f,key_continue)-- The above iterative loop could be mostly parallelized, but would
-- require some care to get the random keys right.sample_brownian_bridgetolerancesamplerkey_drawt0left_ft1right_ft
Suprisinly, a sampler for 1D Brownian motion can easily be extended
to sample from a Brownian sheet of any dimension.
We simply have to replace the sampler used for scalar-valued
standard Gaussian noise with one that samples entire functions
from a Gaussian process. In particular, we can re-use the
1D Brownian bridge sampler that we wrote above.
This process can be nested to arbitrarily many dimensions.
instanceHasBrownianSheet((a,Float),v)given(a,v|VSpace)(HasBrownianSheetav)defsample_brownian_sheet(tol,k,xab)=-- the call to `sample_brownian_sheet` below will recurse
-- until the input is one-dimensional.(xa,xb)=xabsample_a=\k.sample_brownian_sheettolkxasample_brownian_motiontolsample_akxb
Demo: Two-dimensional Brownian sheet
Next, we plot a single sample from a two-dimensional Brownian
sheet with a 3-dimensional (red, green, blue) output.