mesh_tissue/numba_functions.py

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