fixed initialization

examples/sediment_deposit/sediment_deposit.py

 import numpy as np
 import scipy as sp
 import sympy as sm
+import numba as nb
 
 TANK_RATIO = 2
 SEDIMENT_COUNT  = 28000
-SEDIMENT_RADIUS = 1e-3
+SEDIMENT_RADIUS = 0.5e-3
 
 # initialize vorticity
 def init_vorticity(data, coords, component=None):
@@ -22,25 +23,79 @@ def init_sediment(data, coords, nblobs, rblob):
     data   = data[0]
     coords = coords[0]
     X, Y = coords
-    Bx = TANK_RATIO*np.random.rand(nblobs)
-    By = 1*np.random.rand(nblobs)
     R2 = rblob * rblob
     
     cache_file='/tmp/C_init_{}_{}'.format('_'.join(str(x) for x in data.shape),
             str(abs(hash((TANK_RATIO,nblobs,rblob)))))
     try:
+        raise NotImplementedError
         D = np.load(file=cache_file+'.npz')
         print 'Initializing sediments from cache: "{}.npz".'.format(cache_file)
         data[...] = D['data']
     except:
-        # this just explodes the RAM so we iterate per line
+        # this just explodes the RAM
         #data[...] = 0.5 * ((X[:,:,None] - Bx[None,None,:])**2 + (Y[:,:,None] - By[None,None,:])**2 < R2).any(axis=-1)
+        
+        # we could iterate per line but this takes too much time for considered discretizations
+        #data[...] = 0.0
+        #for i in xrange(Y.size):
+            #print '{}/{}'.format(i+1, Y.size)
+            #mask =  ((X[:,:,None] - Bx[None,None,:])**2 + (Y[i,:,None] - By[None,None,:])**2 < R2).any(axis=-1)
+            #data[i,mask[0]] = 0.5
+
+        # we could also iterate on blobs but this takes too much time for considered number of blobs
+        #mask = np.zeros(shape=data.shape, dtype=np.bool_)
+        #mask[...] = 0
+        #for _ in xrange(nblobs):
+            #print 'blob {}/{}'.format(_+1, nblobs)
+            #x, y = TANK_RATIO*np.random.rand(), np.random.rand()
+            #mask |= ((X-x)**2 + (Y-y)**2 < R2)
+        
+        X,   Y   = X.ravel(), Y.ravel()
+        dx,  dy  = X[1]-X[0], Y[1]-Y[0]
+        Nx,  Ny  = X.size, Y.size
+        Rx,  Ry  = 2+int(rblob/dx), 2+int(rblob/dy)
+        assert (rblob>=dx), 'Sediment radius < dx.'
+        assert (rblob>=dy), 'Sediment radius < dy.'
+        
+        Bx = TANK_RATIO*np.random.rand(nblobs)
+        By = 1*np.random.rand(nblobs)
+        Ix = np.floor(Bx/dx).astype(np.int32)
+        Iy = np.floor(By/dy).astype(np.int32)
+        Px = Bx - Ix*dx
+        Py = By - Iy*dy
+        
+        from hysop.tools.numba_utils import make_numba_signature
+        args = (Ix, Iy, Bx, By, data)
+        signature, _ = make_numba_signature(*args)
+         
+        @nb.guvectorize([signature],
+            '(n),(n),(n),(n),(n0,n1)', 
+            target='parallel',
+            nopython=True, cache=True)
+        def iter_blobs(Ix, Iy, Bx, By, data):
+            for k in xrange(nblobs):
+                #print 'blob {}/{}'.format(k+1, nblobs)
+                ix, iy = Ix[k], Iy[k]
+                px, py = Px[k], Py[k]
+                for i in xrange(-Ry, +Ry):
+                    ii = iy+i
+                    if (ii<0) or (ii>=Ny):
+                        continue
+                    dy2 = (py + i*dy)**2
+                    for j in xrange(-Rx, +Rx):
+                        jj = ix+j
+                        if (jj<0) or (jj>=Nx):
+                            continue
+                        dx2 = (px - j*dx)**2
+                        d = dx2 + dy2
+                        if (d<R2):
+                            data[ii,jj] = 0.5
+
         print 'Initializing sediments of radius {} with {} random blobs.'.format(rblob, nblobs)
-        data[...] = 0.0
-        for i in xrange(Y.size):
-            print '{}/{}'.format(i+1, Y.size)
-            mask =  ((X[:,:,None] - Bx[None,None,:])**2 + (Y[i,:,None] - By[None,None,:])**2 < R2).any(axis=-1)
-            data[i,mask[0]] = 0.5
+        data[...]  = 0.0
+        iter_blobs(*args)
+
         # we cache initialization
         np.savez_compressed(file=cache_file, data=data)
         print 'Caching data to "{}.npz".'.format(cache_file)
@@ -270,6 +325,7 @@ def compute(args):
 
     problem = Problem(method=method)
     problem.insert(poisson, 
+                   #diffuse_S,
                    splitting, 
                    dump_fields,
                    compute_mean_fields,
@@ -337,9 +393,10 @@ if __name__=='__main__':
     parser.set_defaults(impl='cl', ndim=2, 
                         npts=(3072+1,TANK_RATIO*3072+1),
                         box_origin=(0.0,), box_length=(1.0,), 
-                        tstart=0.0, tend=100000.0, 
+                        tstart=0.0, tend=10000.0, 
                         dt=1e-6, cfl=0.5, lcfl=0.95,
-                        dump_freq=250)
+                        dump_times=tuple(float(x) for x in range(0,10000,10)),
+                        dump_freq=0)
 
     parser.run(compute)