/
generic_utils.py
169 lines (135 loc) · 4.52 KB
/
generic_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
# -*- coding: utf-8 -*-
# Copyright 2017-2024 The diffsims developers
#
# This file is part of diffsims.
#
# diffsims is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# diffsims is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with diffsims. If not, see <http://www.gnu.org/licenses/>.
"""
Created on 31 Oct 2019
Generic tools for all areas of code.
@author: Rob Tovey
"""
from numpy import isscalar, zeros, array
import numba
__all__ = [
"GLOBAL_BOOL",
"get_grid",
"to_mesh",
]
# Coverage: Cuda code is not tested
try: # pragma: no cover
from numba import cuda
__CUDA = cuda.is_available()
except Exception: # pragma: no cover
cuda = None
__CUDA = False
class GLOBAL_BOOL:
"""
An object which behaves like a bool but can be changed in-place by `set`
or by calling as a function.
"""
def __init__(self, val):
self.val = bool(val)
def __call__(self, val):
self.set(val)
def set(self, val):
self.val = bool(val)
def __bool__(self):
return self.val
def __str__(self):
return str(self.val)
_CUDA = GLOBAL_BOOL(__CUDA)
def get_grid(sz, tpb=None):
dim = len(sz)
if tpb is None:
# Complete guess that feels reasonable
tpb = min(int(512 ** (1 / dim)), 256)
tpb = [tpb] * dim if isscalar(tpb) else list(tpb)
grid = [0] * dim
for i in range(dim):
if tpb[i] > sz[i]:
tpb[i] = sz[i]
grid[i] = 1
else:
while tpb[i] * (sz[i] // tpb[i]) != sz[i]:
tpb[i] -= 1
grid[i] = sz[i] // tpb[i]
return grid, tpb
# Coverage: Numba code does not register when code is run
@numba.njit(parallel=True, fastmath=True, cache=False)
def __toMesh2d(x0, x1, dx0, dx1, out): # pragma: no cover
for i0 in numba.prange(x0.size):
X00 = x0[i0] * dx0[0]
X01 = x0[i0] * dx0[1]
for i1 in range(x1.size):
out[i0, i1, 0] = X00 + x1[i1] * dx1[0]
out[i0, i1, 1] = X01 + x1[i1] * dx1[1]
# Coverage: Numba code does not register when code is run
@numba.njit(parallel=True, fastmath=True, cache=False)
def __toMesh3d(x0, x1, x2, dx0, dx1, dx2, out): # pragma: no cover
for i0 in numba.prange(x0.size):
X00 = x0[i0] * dx0[0]
X01 = x0[i0] * dx0[1]
X02 = x0[i0] * dx0[2]
for i1 in range(x1.size):
X10 = x1[i1] * dx1[0]
X11 = x1[i1] * dx1[1]
X12 = x1[i1] * dx1[2]
for i2 in range(x2.size):
out[i0, i1, i2, 0] = X00 + X10 + x2[i2] * dx2[0]
out[i0, i1, i2, 1] = X01 + X11 + x2[i2] * dx2[1]
out[i0, i1, i2, 2] = X02 + X12 + x2[i2] * dx2[2]
def to_mesh(x, dx=None, dtype=None):
"""
Generates dense meshes from grid vectors, broadly:
to_mesh(x)[i,j,...] = (x[0][i], x[1][j], ...)
Parameters
----------
x : `list` [`numpy.ndarray`], of shape [(nx,), (ny,), ...]
List of grid vectors
dx : `list` [`numpy.ndarray`] or `None`, optional
Basis in which to expand mesh, default is the canonical basis
dtype : `str` or `None`, optional
String representing the `numpy` type of output, default inherits from `x`
Returns
-------
X : `numpy.ndarray` [dtype], (x[0].size, x[1].size, ..., len(x))
X[i,j,..., k] = x[0][i]*dx[0][k] + x[1][j]*dx[1][k] + ...
"""
shape = [xi.size for xi in x]
if dtype is None:
dtype = x[0].dtype
else:
x = [xi.astype(dtype, copy=False) for xi in x]
dim = len(shape)
X = zeros(shape + [dim], dtype=dtype)
if dim == 2:
dx = [array([1, 0]), array([0, 1])] if dx is None else list(dx)
__toMesh2d(*x, *dx, X)
return X
elif dim == 3:
dx = (
[array([1, 0, 0]), array([0, 1, 0]), array([0, 0, 1])]
if dx is None
else list(dx)
)
__toMesh3d(*x, *dx, X)
return X
# TODO: this ignores dx, probably just raise error on dim=4?
dim = X.shape[-1]
for i in range(len(x)):
sz = [1] * dim
sz[i] = -1
X[..., i] += x[i].reshape(sz)
return X