mesh_tissue/numba_functions.py

   1 # SPDX-FileCopyrightText: 2019-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import numpy as np
   6 import time
   7 import sys
   8
   9 bool_numba = False
  10
  11 try:
  12     from .utils_pip import Pip
  13     Pip._ensure_user_site_package()
  14     from numba import jit, njit, guvectorize, float64, int32, prange
  15     from numba.typed import List
  16     bool_numba = True
  17 except:
  18     pass
  19     '''
  20     try:
  21         from .utils_pip import Pip
  22         #Pip.upgrade_pip()
  23         Pip.install('llvmlite')
  24         Pip.install('numba')
  25         from numba import jit, njit, guvectorize, float64, int32, prange
  26         bool_numba = True
  27         print('Tissue: Numba successfully installed!')
  28     except:
  29         print('Tissue: Numba not loaded correctly. Try restarting Blender')
  30     '''
  31
  32 if bool_numba:
  33     #from numba import jit, njit, guvectorize, float64, int32, prange
  34
  35
  36     @njit(parallel=True)
  37     #@cuda.jit('void(float32[:,:], float32[:,:])')
  38     def tex_laplacian(lap, arr):
  39         arr2 = arr*2
  40         diag = sqrt(2)/2
  41         nx = arr.shape[0]
  42         ny = arr.shape[1]
  43         for i in prange(nx):
  44             for j in prange(ny):
  45                 i0 = (i-1)%nx
  46                 j0 = (j-1)%ny
  47                 i1 = (i+1)%nx
  48                 j1 = (j+1)%ny
  49                 #lap[i+1,j+1] = arr[i, j+1] + arr[i+2, j+1] + arr[i+1, j] + arr[i+1, j+2] - 4*arr[i+1,j+1]
  50
  51                 lap[i,j] = ((arr[i0, j] + arr[i1, j] - arr2[i,j]) + \
  52                                (arr[i, j0] + arr[i, j1] - arr2[i,j]) + \
  53                                (arr[i0, j0] + arr[i1, j1] - arr2[i,j])*diag + \
  54                                (arr[i1, j0] + arr[i0, j1] - arr2[i,j])*diag)*0.75
  55
  56     @njit(parallel=True)
  57     def tex_laplacian_ani(lap, arr, VF):
  58         arr2 = arr*2
  59         nx = arr.shape[0]
  60         ny = arr.shape[1]
  61         i0 = np.arange(nx)-1
  62         i0[0] = 1
  63         i1 = np.arange(nx)+1
  64         i1[nx-1] = nx-2
  65         j0 = np.arange(ny)-1
  66         j0[0] = 1
  67         j1 = np.arange(ny)+1
  68         j1[ny-1] = ny-2
  69         for i in prange(nx):
  70             for j in prange(ny):
  71                 lap[i,j] = (arr[i0[i], j] + arr[i1[i], j] - arr2[i,j])*VF[0,i,j] + \
  72                                (arr[i, j0[j]] + arr[i, j1[j]] - arr2[i,j])*VF[1,i,j] + \
  73                                (arr[i0[i], j0[j]] + arr[i1[i], j1[j]] - arr2[i,j])*VF[2,i,j] + \
  74                                (arr[i1[i], j0[j]] + arr[i0[i], j1[j]] - arr2[i,j])*VF[3,i,j]
  75         #lap[0,:] = lap[1,:]
  76         #lap[:,0] = lap[:,1]
  77         #lap[-1,:] = lap[-2,:]
  78         #lap[:,-1] = lap[:,-2]
  79
  80     #@cuda.jit(parallel=True)
  81     @njit(parallel=True)
  82     def run_tex_rd(A, B, lap_A, lap_B, diff_A, diff_B, f, k, dt, steps, brush):
  83         for t in range(steps):
  84             tex_laplacian(lap_A, A)
  85             tex_laplacian(lap_B, B)
  86             nx = A.shape[0]
  87             ny = A.shape[1]
  88             for i in prange(nx):
  89                 for j in prange(ny):
  90                     B[i,j] += brush[i,j]
  91                     ab2 =  A[i,j]*B[i,j]**2
  92                     A[i,j] += (lap_A[i,j]*diff_A - ab2 + f*(1-A[i,j]))*dt
  93                     B[i,j] += (lap_B[i,j]*diff_B + ab2 - (k+f)*B[i,j])*dt
  94
  95     @njit(parallel=True)
  96     def run_tex_rd_ani(A, B, lap_A, lap_B, diff_A, diff_B, f, k, dt, steps, vf1, vf2, brush):
  97         for t in range(steps):
  98             tex_laplacian_ani(lap_A, A, vf2)
  99             #laplacian(lap_A, A)
 100             tex_laplacian_ani(lap_B, B, vf1)
 101             nx = A.shape[0]
 102             ny = A.shape[1]
 103             for i in prange(nx):
 104                 for j in prange(ny):
 105                     B[i,j] += brush[i,j]
 106                     ab2 =  A[i ,j]*B[i,j]**2
 107                     A[i,j] += (lap_A[i,j]*diff_A[i,j] - ab2 + f[i,j]*(1-A[i,j]))*dt
 108                     B[i,j] += (lap_B[i,j]*diff_B[i,j] + ab2 - (k[i,j]+f[i,j])*B[i,j])*dt
 109
 110
 111     @njit(parallel=True)
 112     def numba_reaction_diffusion(n_verts, n_edges, edge_verts, a, b, brush, diff_a, diff_b, f, k, dt, time_steps):
 113         arr = np.arange(n_edges)
 114         id0 = edge_verts[arr*2]
 115         id1 = edge_verts[arr*2+1]
 116         for i in range(time_steps):
 117             lap_a, lap_b = rd_init_laplacian(n_verts)
 118             numba_rd_laplacian(id0, id1, a, b, lap_a, lap_b)
 119             numba_rd_core(a, b, lap_a, lap_b, diff_a, diff_b, f, k, dt)
 120             numba_set_ab(a,b,brush)
 121         return a,b
 122
 123     @njit(parallel=False)
 124     def integrate_field(n_edges, id0, id1, values, edge_flow, mult, time_steps):
 125         #n_edges = len(edge_flow)
 126         for i in range(time_steps):
 127             values0 = values
 128             for j in range(n_edges):
 129                 v0 = id0[j]
 130                 v1 = id1[j]
 131                 values[v0] -= values0[v1] * edge_flow[j] * 0.001#mult[v1]
 132                 values[v1] += values0[v0] * edge_flow[j] * 0.001#mult[v0]
 133             for j in range(n_edges):
 134                 v0 = id0[j]
 135                 v1 = id1[j]
 136                 values[v0] = max(values[v0],0)
 137                 values[v1] = max(values[v1],0)
 138         return values
 139
 140     @njit(parallel=True)
 141     def numba_reaction_diffusion_anisotropic(n_verts, n_edges, edge_verts, a, b, brush, diff_a, diff_b, f, k, dt, time_steps, field_mult):
 142         arr = np.arange(n_edges)
 143         id0 = edge_verts[arr*2]
 144         id1 = edge_verts[arr*2+1]
 145         mult = field_mult[arr]
 146         for i in range(time_steps):
 147             lap_a, lap_b = rd_init_laplacian(n_verts)
 148             numba_rd_laplacian_anisotropic(id0, id1, a, b, lap_a, lap_b, mult)
 149             numba_rd_core(a, b, lap_a, lap_b, diff_a, diff_b, f, k, dt)
 150             numba_set_ab(a,b,brush)
 151         return a,b
 152
 153     #@guvectorize(['(float64[:] ,float64[:] , float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64)'],'(n),(n),(n),(n),(n),(n),(n),(n),()',target='parallel')
 154     @njit(parallel=True)
 155     def numba_rd_core(a, b, lap_a, lap_b, diff_a, diff_b, f, k, dt):
 156         n = len(a)
 157         _f = np.full(n, f[0]) if len(f) == 1 else f
 158         _k = np.full(n, k[0]) if len(k) == 1 else k
 159         _diff_a = np.full(n, diff_a[0]) if len(diff_a) == 1 else diff_a
 160         _diff_b = np.full(n, diff_b[0]) if len(diff_b) == 1 else diff_b
 161
 162         for i in prange(n):
 163             fi = _f[i]
 164             ki = _k[i]
 165             diff_ai = _diff_a[i]
 166             diff_bi = _diff_b[i]
 167             ab2 = a[i]*b[i]**2
 168             a[i] += (diff_ai * lap_a[i] - ab2 + fi*(1-a[i]))*dt
 169             b[i] += (diff_bi * lap_b[i] + ab2 - (ki+fi)*b[i])*dt
 170
 171     @njit(parallel=True)
 172     def numba_rd_core_(a, b, lap_a, lap_b, diff_a, diff_b, f, k, dt):
 173         ab2 = a*b**2
 174         a += (diff_a*lap_a - ab2 + f*(1-a))*dt
 175         b += (diff_b*lap_b + ab2 - (k+f)*b)*dt
 176
 177     @njit(parallel=True)
 178     def numba_set_ab(a, b, brush):
 179         n = len(a)
 180         _brush = np.full(n, brush[0]) if len(brush) == 1 else brush
 181         for i in prange(len(b)):
 182             b[i] += _brush[i]
 183             if b[i] < 0: b[i] = 0
 184             elif b[i] > 1: b[i] = 1
 185             if a[i] < 0: a[i] = 0
 186             elif a[i] > 1: a[i] = 1
 187
 188     @njit(parallel=True)
 189     def numba_rd_laplacian(id0, id1, a, b, lap_a, lap_b):
 190         for i in prange(len(id0)):
 191             v0 = id0[i]
 192             v1 = id1[i]
 193             lap_a[v0] += (a[v1] - a[v0])
 194             lap_a[v1] += (a[v0] - a[v1])
 195             lap_b[v0] += (b[v1] - b[v0])
 196             lap_b[v1] += (b[v0] - b[v1])
 197         #return lap_a, lap_b
 198
 199     @njit(parallel=True)
 200     def numba_rd_laplacian_anisotropic(id0, id1, a, b, lap_a, lap_b, mult):
 201         for i in prange(len(id0)):
 202             v0 = id0[i]
 203             v1 = id1[i]
 204             multiplier = mult[i]
 205             lap_a[v0] += (a[v1] - a[v0])# * multiplier
 206             lap_a[v1] += (a[v0] - a[v1])# * multiplier
 207             lap_b[v0] += (b[v1] - b[v0]) * multiplier
 208             lap_b[v1] += (b[v0] - b[v1]) * multiplier
 209         #return lap_a, lap_b
 210
 211     @njit(parallel=True)
 212     def numba_rd_neigh_vertices(edge_verts):
 213         n_edges = len(edge_verts)/2
 214         id0 = np.zeros(n_edges)
 215         id1 = np.zeros(n_edges)
 216         for i in prange(n_edges):
 217             id0[i] = edge_verts[i*2]     # first vertex indices for each edge
 218             id1[i] = edge_verts[i*2+1]   # second vertex indices for each edge
 219         return id0, id1
 220
 221     #@guvectorize(['(float64[:] ,float64[:] , float64[:], float64[:], float64[:])'],'(m),(n),(n),(n),(n)',target='parallel')
 222     @njit(parallel=True)
 223     #@njit
 224     def numba_rd_laplacian_(edge_verts, a, b, lap_a, lap_b):
 225         for i in prange(len(edge_verts)/2):
 226             v0 = edge_verts[i*2]
 227             v1 = edge_verts[i*2+1]
 228             lap_a[v0] += a[v1] - a[v0]
 229             lap_a[v1] += a[v0] - a[v1]
 230             lap_b[v0] += b[v1] - b[v0]
 231             lap_b[v1] += b[v0] - b[v1]
 232         #return lap_a, lap_b
 233
 234     @njit(parallel=True)
 235     def rd_fill_laplacian(lap_a, lap_b, id0, id1, lap_a0, lap_b0):
 236         #for i, j, la0, lb0 in zip(id0,id1,lap_a0,lap_b0):
 237         for index in prange(len(id0)):
 238             i = id0[index]
 239             j = id1[index]
 240             la0 = lap_a0[index]
 241             lb0 = lap_b0[index]
 242             lap_a[i] += la0
 243             lap_b[i] += lb0
 244             lap_a[j] -= la0
 245             lap_b[j] -= lb0
 246
 247     @njit(parallel=True)
 248     def rd_init_laplacian(n_verts):
 249         lap_a = np.zeros(n_verts)
 250         lap_b = np.zeros(n_verts)
 251         return lap_a, lap_b
 252
 253     '''
 254     @jit
 255     def numba_reaction_diffusion(n_verts, n_edges, edge_verts, a, b, diff_a, diff_b, f, k, dt, time_steps, db):
 256         arr = np.arange(n_edges)*2
 257         id0 = edge_verts[arr]     # first vertex indices for each edge
 258         id1 = edge_verts[arr+1]   # second vertex indices for each edge
 259         #dgrad = abs(grad[id1] - grad[id0])
 260         for i in range(time_steps):
 261             lap_a = np.zeros(n_verts)
 262             lap_b = np.zeros(n_verts)
 263             b += db
 264             lap_a0 =  a[id1] - a[id0]   # laplacian increment for first vertex of each edge
 265             lap_b0 =  b[id1] -  b[id0]   # laplacian increment for first vertex of each edge
 266             #lap_a0 *= dgrad
 267             #lap_b0 *= dgrad
 268
 269             for i, j, la0, lb0 in zip(id0,id1,lap_a0,lap_b0):
 270                 lap_a[i] += la0
 271                 lap_b[i] += lb0
 272                 lap_a[j] -= la0
 273                 lap_b[j] -= lb0
 274             ab2 = a*b**2
 275             #a += eval("(diff_a*lap_a - ab2 + f*(1-a))*dt")
 276             #b += eval("(diff_b*lap_b + ab2 - (k+f)*b)*dt")
 277             a += (diff_a*lap_a - ab2 + f*(1-a))*dt
 278             b += (diff_b*lap_b + ab2 - (k+f)*b)*dt
 279         return a, b
 280     '''
 281     '''
 282     @njit(parallel=True)
 283     def numba_lerp2_(v00, v10, v01, v11, vx, vy):
 284         sh = v00.shape
 285         co2 = np.zeros((sh[0],len(vx),sh[-1]))
 286         for i in prange(len(v00)):
 287             for j in prange(len(vx)):
 288                 for k in prange(len(v00[0][0])):
 289                     co0 = v00[i][0][k] + (v10[i][0][k] - v00[i][0][k]) * vx[j][0]
 290                     co1 = v01[i][0][k] + (v11[i][0][k] - v01[i][0][k]) * vx[j][0]
 291                     co2[i][j][k] = co0 + (co1 - co0) * vy[j][0]
 292         return co2
 293
 294
 295     @njit(parallel=True)
 296     def numba_lerp2_vec(v0, vx, vy):
 297         n_faces = v0.shape[0]
 298         co2 = np.zeros((n_faces,len(vx),3))
 299         for i in prange(n_faces):
 300             for j in prange(len(vx)):
 301                 for k in prange(3):
 302                     co0 = v0[i][0][k] + (v0[i][1][k] - v0[i][0][k]) * vx[j][0]
 303                     co1 = v0[i][3][k] + (v0[i][2][k] - v0[i][3][k]) * vx[j][0]
 304                     co2[i][j][k] = co0 + (co1 - co0) * vy[j][0]
 305         return co2
 306
 307     @njit(parallel=True)
 308     def numba_lerp2__(val, vx, vy):
 309         n_faces = len(val)
 310         co2 = np.zeros((n_faces,len(vx),1))
 311         for i in prange(n_faces):
 312             for j in prange(len(vx)):
 313                 co0 = val[i][0] + (val[i][1] - val[i][0]) * val[j][0]
 314                 co1 = val[i][3] + (val[i][2] - val[i][3]) * val[j][0]
 315                 co2[i][j][0] = co0 + (co1 - co0) * vy[j][0]
 316         return co2
 317     '''
 318
 319     @njit(parallel=True)
 320     def numba_combine_and_flatten(arrays):
 321         n_faces = len(arrays)
 322         n_verts = len(arrays[0])
 323         new_list = [0.0]*n_faces*n_verts*3
 324         for i in prange(n_faces):
 325             for j in prange(n_verts):
 326                 for k in prange(3):
 327                     new_list[i*n_verts*3+j*3+k] = arrays[i][j,k]
 328         return new_list
 329
 330     @njit(parallel=True)
 331     def numba_calc_thickness_area_weight(co2,n2,vz,a,weight):
 332         shape = co2.shape
 333         n_patches = shape[0]
 334         n_verts = shape[1]
 335         n_co = shape[2]
 336         nn = n2.shape[1]-1
 337         na = a.shape[1]-1
 338         nw = weight.shape[1]-1
 339         co3 = np.zeros((n_patches,n_verts,n_co))
 340         for i in prange(n_patches):
 341             for j in prange(n_verts):
 342                 for k in prange(n_co):
 343                     co3[i,j,k] = co2[i,j,k] + n2[i,min(j,nn),k] * vz[0,j,0] * a[i,min(j,na),0] * weight[i,min(j,nw),0]
 344         return co3
 345     '''
 346     @njit(parallel=True)
 347     def numba_calc_thickness_area(co2,n2,vz,a):
 348         shape = co2.shape
 349         n_patches = shape[0]
 350         n_verts = shape[1]
 351         n_co = shape[2]
 352         #co3 = [0.0]*n_patches*n_verts*n_co #np.zeros((n_patches,n_verts,n_co))
 353         co3 = np.zeros((n_patches,n_verts,n_co))
 354         for i in prange(n_patches):
 355             for j in prange(n_verts):
 356                 for k in prange(n_co):
 357                     #co3[i,j,k] = co2[i,j,k] + n2[i,j,k] * vz[0,j,0] * a[i,j,0]
 358                     co3[i,j,k] = co2[i,j,k] + n2[i,min(j,nor_len),k] * vz[0,j,0] * a[i,j,0]
 359         return co3
 360     '''
 361     @njit(parallel=True)
 362     def numba_calc_thickness_weight(co2,n2,vz,weight):
 363         shape = co2.shape
 364         n_patches = shape[0]
 365         n_verts = shape[1]
 366         n_co = shape[2]
 367         nn = n2.shape[1]-1
 368         nw = weight.shape[1]-1
 369         co3 = np.zeros((n_patches,n_verts,n_co))
 370         for i in prange(n_patches):
 371             for j in prange(n_verts):
 372                 for k in prange(n_co):
 373                     co3[i,j,k] = co2[i,j,k] + n2[i,min(j,nn),k] * vz[0,j,0] * weight[i,min(j,nw),0]
 374         return co3
 375
 376     @njit(parallel=True)
 377     def numba_calc_thickness(co2,n2,vz):
 378         shape = co2.shape
 379         n_patches = shape[0]
 380         n_verts = shape[1]
 381         n_co = shape[2]
 382         nn = n2.shape[1]-1
 383         co3 = np.zeros((n_patches,n_verts,n_co))
 384         for i in prange(n_patches):
 385             for j in prange(n_verts):
 386                 for k in prange(n_co):
 387                     co3[i,j,k] = co2[i,j,k] + n2[i,min(j,nn),k] * vz[0,j,0]
 388         return co3
 389
 390     @njit(parallel=True)
 391     def numba_interp_points(v00, v10, v01, v11, vx, vy):
 392         n_patches = v00.shape[0]
 393         n_verts = vx.shape[1]
 394         n_verts0 = v00.shape[1]
 395         n_co = v00.shape[2]
 396         vxy = np.zeros((n_patches,n_verts,n_co))
 397         for i in prange(n_patches):
 398             for j in prange(n_verts):
 399                 j0 = min(j,n_verts0-1)
 400                 for k in prange(n_co):
 401                     co0 = v00[i,j0,k] + (v10[i,j0,k] - v00[i,j0,k]) * vx[0,j,0]
 402                     co1 = v01[i,j0,k] + (v11[i,j0,k] - v01[i,j0,k]) * vx[0,j,0]
 403                     vxy[i,j,k] = co0 + (co1 - co0) * vy[0,j,0]
 404         return vxy
 405
 406     @njit(parallel=True)
 407     def numba_interp_points_sk(v00, v10, v01, v11, vx, vy):
 408         n_patches = v00.shape[0]
 409         n_sk = v00.shape[1]
 410         n_verts = v00.shape[2]
 411         n_co = v00.shape[3]
 412         vxy = np.zeros((n_patches,n_sk,n_verts,n_co))
 413         for i in prange(n_patches):
 414             for sk in prange(n_sk):
 415                 for j in prange(n_verts):
 416                     for k in prange(n_co):
 417                         co0 = v00[i,sk,j,k] + (v10[i,sk,j,k] - v00[i,sk,j,k]) * vx[0,sk,j,0]
 418                         co1 = v01[i,sk,j,k] + (v11[i,sk,j,k] - v01[i,sk,j,k]) * vx[0,sk,j,0]
 419                         vxy[i,sk,j,k] = co0 + (co1 - co0) * vy[0,sk,j,0]
 420         return vxy
 421
 422     @njit
 423     def numba_lerp(v0, v1, x):
 424         return v0 + (v1 - v0) * x
 425
 426     @njit
 427     def numba_lerp2(v00, v10, v01, v11, vx, vy):
 428         co0 = numba_lerp(v00, v10, vx)
 429         co1 = numba_lerp(v01, v11, vx)
 430         co2 = numba_lerp(co0, co1, vy)
 431         return co2
 432
 433     @njit(parallel=True)
 434     def numba_lerp2_________________(v00, v10, v01, v11, vx, vy):
 435         ni = len(v00)
 436         nj = len(v00[0])
 437         nk = len(v00[0][0])
 438         co2 = np.zeros((ni,nj,nk))
 439         for i in prange(ni):
 440             for j in prange(nj):
 441                 for k in prange(nk):
 442                     _v00 = v00[i,j,k]
 443                     _v01 = v01[i,j,k]
 444                     _v10 = v10[i,j,k]
 445                     _v11 = v11[i,j,k]
 446                     co0 = _v00 + (_v10 - _v00) * vx[i,j,k]
 447                     co1 = _v01 + (_v11 - _v01) * vx[i,j,k]
 448                     co2[i,j,k] = co0 + (co1 - co0) * vy[i,j,k]
 449         return co2
 450
 451     @njit(parallel=True)
 452     def numba_lerp2_4(v00, v10, v01, v11, vx, vy):
 453         ni = len(v00)
 454         nj = len(v00[0])
 455         nk = len(v00[0][0])
 456         nw = len(v00[0][0][0])
 457         co2 = np.zeros((ni,nj,nk,nw))
 458         for i in prange(ni):
 459             for j in prange(nj):
 460                 for k in prange(nk):
 461                     for w in prange(nw):
 462                         _v00 = v00[i,j,k]
 463                         _v01 = v01[i,j,k]
 464                         _v10 = v10[i,j,k]
 465                         _v11 = v11[i,j,k]
 466                         co0 = _v00 + (_v10 - _v00) * vx[i,j,k]
 467                         co1 = _v01 + (_v11 - _v01) * vx[i,j,k]
 468                         co2[i,j,k] = co0 + (co1 - co0) * vy[i,j,k]
 469         return co2
 470
 471
 472 #except:
 473 #    print("Tissue: Numba cannot be installed. Try to restart Blender.")
 474 #    pass