spinn 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- spinn-0.1.0/LICENSE +21 -0
- spinn-0.1.0/PKG-INFO +23 -0
- spinn-0.1.0/README.md +2 -0
- spinn-0.1.0/pyproject.toml +3 -0
- spinn-0.1.0/setup.cfg +4 -0
- spinn-0.1.0/setup.py +25 -0
- spinn-0.1.0/spinn/__init__.py +6 -0
- spinn-0.1.0/spinn/graphs.py +93 -0
- spinn-0.1.0/spinn/trees.py +295 -0
- spinn-0.1.0/spinn/utils.py +166 -0
- spinn-0.1.0/spinn.egg-info/PKG-INFO +23 -0
- spinn-0.1.0/spinn.egg-info/SOURCES.txt +13 -0
- spinn-0.1.0/spinn.egg-info/dependency_links.txt +1 -0
- spinn-0.1.0/spinn.egg-info/requires.txt +5 -0
- spinn-0.1.0/spinn.egg-info/top_level.txt +1 -0
spinn-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2024 Marek Wojciechowski
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
spinn-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: spinn
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Simple Physics Informed Neural Networks
|
|
5
|
+
Home-page: https://github.com/mrkwjc/spinn
|
|
6
|
+
Author: Marek Wojciechowski
|
|
7
|
+
Author-email: mwojc@p.lodz.pl
|
|
8
|
+
License: LGPL-2.1
|
|
9
|
+
Keywords: ann,pinn
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Operating System :: OS Independent
|
|
13
|
+
Requires-Python: >=3.6
|
|
14
|
+
Description-Content-Type: text/markdown
|
|
15
|
+
License-File: LICENSE
|
|
16
|
+
Requires-Dist: jax
|
|
17
|
+
Requires-Dist: jaxlib
|
|
18
|
+
Requires-Dist: numpy
|
|
19
|
+
Requires-Dist: scipy
|
|
20
|
+
Requires-Dist: networkx
|
|
21
|
+
|
|
22
|
+
# spinn - Simple Physics Informed Neural Networks
|
|
23
|
+
Makes PINNs easy (with jax), see examples...
|
spinn-0.1.0/README.md
ADDED
spinn-0.1.0/setup.cfg
ADDED
spinn-0.1.0/setup.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
import setuptools
|
|
2
|
+
|
|
3
|
+
with open("README.md", "r") as fh:
|
|
4
|
+
long_description = fh.read()
|
|
5
|
+
|
|
6
|
+
setuptools.setup(
|
|
7
|
+
name="spinn",
|
|
8
|
+
version="0.1.0",
|
|
9
|
+
author="Marek Wojciechowski",
|
|
10
|
+
author_email="mwojc@p.lodz.pl",
|
|
11
|
+
description="Simple Physics Informed Neural Networks",
|
|
12
|
+
long_description=long_description,
|
|
13
|
+
long_description_content_type="text/markdown",
|
|
14
|
+
url="https://github.com/mrkwjc/spinn",
|
|
15
|
+
packages=setuptools.find_packages(),
|
|
16
|
+
keywords=['ann', 'pinn'],
|
|
17
|
+
license='LGPL-2.1',
|
|
18
|
+
classifiers=[
|
|
19
|
+
"Programming Language :: Python :: 3",
|
|
20
|
+
"License :: OSI Approved :: MIT License",
|
|
21
|
+
"Operating System :: OS Independent",
|
|
22
|
+
],
|
|
23
|
+
python_requires='>=3.6',
|
|
24
|
+
install_requires=['jax', 'jaxlib', 'numpy', 'scipy', 'networkx'],
|
|
25
|
+
)
|
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import networkx as nx
|
|
3
|
+
from spinn import utils
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def _shoot(src, trg, mindeg, d=None, reverse=False):
|
|
7
|
+
mindeg = min(mindeg, len(trg))
|
|
8
|
+
if d is None:
|
|
9
|
+
s = np.repeat(src, mindeg)
|
|
10
|
+
t = [np.random.choice(trg, mindeg, replace=False)
|
|
11
|
+
for i in range(len(src))]
|
|
12
|
+
else:
|
|
13
|
+
s, t = [], []
|
|
14
|
+
for i, ei in enumerate(d):
|
|
15
|
+
toshoot = mindeg-len(ei)
|
|
16
|
+
if toshoot > 0:
|
|
17
|
+
s += [src[i]]*toshoot
|
|
18
|
+
if toshoot == mindeg:
|
|
19
|
+
ti = trg
|
|
20
|
+
else:
|
|
21
|
+
ei = np.asarray(ei)
|
|
22
|
+
ni = ei[:, 0] if reverse else ei[:, 1]
|
|
23
|
+
ti = np.setdiff1d(trg, ni)
|
|
24
|
+
t += [np.random.choice(ti, toshoot, replace=False)]
|
|
25
|
+
if not len(s):
|
|
26
|
+
return []
|
|
27
|
+
e = np.asarray((s, np.concatenate(t)))
|
|
28
|
+
e = e[::-1] if reverse else e
|
|
29
|
+
return e.T.tolist()
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def _multilayer(arch, full=False, mindeg=None, G=None):
|
|
33
|
+
extents = utils.pacc((0,) + tuple(arch))
|
|
34
|
+
layers = [range(start, end) for start, end in extents]
|
|
35
|
+
G = nx.DiGraph() if G is None else G
|
|
36
|
+
for i, layer in enumerate(layers):
|
|
37
|
+
G.add_nodes_from(layer, layer=i)
|
|
38
|
+
if not full:
|
|
39
|
+
for layer1, layer2 in utils.pairwise(layers):
|
|
40
|
+
if mindeg is None:
|
|
41
|
+
G.add_edges_from(utils.prod(layer1, layer2))
|
|
42
|
+
else:
|
|
43
|
+
e = _shoot(layer1, layer2, mindeg)
|
|
44
|
+
G.add_edges_from(e)
|
|
45
|
+
d = [tuple(G.in_edges(n)) for n in layer2]
|
|
46
|
+
e = _shoot(layer2, layer1, mindeg, d, reverse=True)
|
|
47
|
+
G.add_edges_from(e)
|
|
48
|
+
else:
|
|
49
|
+
if mindeg is None:
|
|
50
|
+
for i, layer1 in enumerate(layers[:-1]):
|
|
51
|
+
for layer2 in layers[i+1:]:
|
|
52
|
+
G.add_edges_from(utils.prod(layer1, layer2))
|
|
53
|
+
else:
|
|
54
|
+
for i, layer1 in enumerate(layers[:-1]):
|
|
55
|
+
trg = np.concatenate(layers[i+1:])
|
|
56
|
+
e = _shoot(layer1, trg, mindeg)
|
|
57
|
+
G.add_edges_from(e)
|
|
58
|
+
for i, layer2 in enumerate(layers[::-1][:-1]):
|
|
59
|
+
trg = np.concatenate(layers[::-1][i+1:])
|
|
60
|
+
d = [tuple(G.in_edges(n)) for n in layer2]
|
|
61
|
+
e = _shoot(layer2, trg, mindeg, d, reverse=True)
|
|
62
|
+
G.add_edges_from(e)
|
|
63
|
+
return G
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
class mlgraph(nx.DiGraph):
|
|
67
|
+
def __init__(self, arch=None, full=False, d=None):
|
|
68
|
+
nx.DiGraph.__init__(self)
|
|
69
|
+
if arch is not None:
|
|
70
|
+
_multilayer(arch, full, d, self)
|
|
71
|
+
|
|
72
|
+
def preview(self, fname=None, **kwargs):
|
|
73
|
+
pos = nx.multipartite_layout(self, subset_key="layer")
|
|
74
|
+
import matplotlib.pyplot as plt
|
|
75
|
+
plt.figure(figsize=(8, 4))
|
|
76
|
+
nx.draw_networkx(self, pos, **kwargs)
|
|
77
|
+
# plt.axis("equal")
|
|
78
|
+
if fname is not None:
|
|
79
|
+
plt.savefig(fname, dpi=300, bbox_inches="tight")
|
|
80
|
+
plt.show()
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
if __name__ == '__main__':
|
|
84
|
+
g = mlgraph((5, 20, 20, 1), full=True, d=3)
|
|
85
|
+
# g.preview()
|
|
86
|
+
g.preview(fname=None,
|
|
87
|
+
with_labels=False,
|
|
88
|
+
node_size=50,
|
|
89
|
+
node_color='k',
|
|
90
|
+
edge_color='grey',
|
|
91
|
+
width=0.5,
|
|
92
|
+
linewidths=0.5,
|
|
93
|
+
arrows=False)
|
|
@@ -0,0 +1,295 @@
|
|
|
1
|
+
import networkx as nx
|
|
2
|
+
import numpy as np
|
|
3
|
+
import jax
|
|
4
|
+
import jax.numpy as jnp
|
|
5
|
+
from functools import partial
|
|
6
|
+
from scipy.optimize import minimize
|
|
7
|
+
from jax.experimental.sparse import BCOO
|
|
8
|
+
from spinn import utils, graphs
|
|
9
|
+
|
|
10
|
+
# import os
|
|
11
|
+
# os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
|
12
|
+
# jax.config.update("jax_enable_x64", True)
|
|
13
|
+
# jax.config.update('jax_default_matmul_precision', 'high') # 'bfloat16_3x'
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class spinn:
|
|
17
|
+
def __init__(self, G, f=None, p=None,
|
|
18
|
+
initx=None, inity=None,
|
|
19
|
+
key=None,
|
|
20
|
+
sparse=False):
|
|
21
|
+
if G is None:
|
|
22
|
+
return # for tree unflattenning
|
|
23
|
+
if isinstance(G, (list, tuple, np.ndarray)):
|
|
24
|
+
G = graphs.mlgraph(G)
|
|
25
|
+
assert nx.is_directed_acyclic_graph(G)
|
|
26
|
+
layers = [sorted(g) for g in nx.topological_generations(G)]
|
|
27
|
+
nlst = sum(layers, [])
|
|
28
|
+
H = nx.relabel_nodes(G, dict(zip(nlst, range(len(G)))))
|
|
29
|
+
layers = [sorted(g) for g in nx.topological_generations(H)]
|
|
30
|
+
arch = [len(li) for li in layers]
|
|
31
|
+
carch = utils.acc(arch) # list(itertools.accumulate(arch))
|
|
32
|
+
sizes = list(zip(arch[1:], carch[:-1])) # sizes of weight matrices
|
|
33
|
+
conecs = [list(zip(*H.in_edges(li))) for li in layers[1:]]
|
|
34
|
+
wmrk = utils.pacc([0] + [len(ci[0]) for ci in conecs])
|
|
35
|
+
bmrk = utils.pacc([wmrk[-1][1]] + arch[1:])
|
|
36
|
+
inps = sorted((n for n, d in G.in_degree() if d == 0))
|
|
37
|
+
outs = sorted((n for n, d in G.out_degree() if d == 0))
|
|
38
|
+
f, p = utils.unroll_fp(f, p, arch)
|
|
39
|
+
self.nlst = nlst
|
|
40
|
+
self.n = len(H)
|
|
41
|
+
self.inps = inps
|
|
42
|
+
self.outs = outs
|
|
43
|
+
self.ni = len(self.inps)
|
|
44
|
+
self.no = len(self.outs)
|
|
45
|
+
self.nw = len(H.edges)
|
|
46
|
+
self.nb = self.n - self.ni
|
|
47
|
+
self.wshp = [(s, (tuple(np.asarray(i1)-i1[0]), i0), wm, bm)
|
|
48
|
+
for s, (i0, i1), wm, bm in zip(sizes, conecs, wmrk, bmrk)]
|
|
49
|
+
self.f = f
|
|
50
|
+
self.p = p
|
|
51
|
+
self.sparse = sparse
|
|
52
|
+
self.initw(key=key)
|
|
53
|
+
if initx is not None:
|
|
54
|
+
self.initx(initx)
|
|
55
|
+
if inity is not None:
|
|
56
|
+
self.inity(inity)
|
|
57
|
+
# self.codes = {None: [[i] for i in range(self.ni)],
|
|
58
|
+
# 'F': [[0]], 'dFx': [[0, 0]]}
|
|
59
|
+
self._register_pytree()
|
|
60
|
+
|
|
61
|
+
def initw(self, key=None):
|
|
62
|
+
w = []
|
|
63
|
+
b = []
|
|
64
|
+
for s, (i0, i1), _, _ in self.wshp:
|
|
65
|
+
u, c = np.unique(i0, return_counts=True)
|
|
66
|
+
n = np.repeat(c, c)
|
|
67
|
+
state = np.random.get_state()
|
|
68
|
+
np.random.seed(key)
|
|
69
|
+
r = np.random.uniform(-1., 1., size=len(i0)+s[0])
|
|
70
|
+
np.random.set_state(state)
|
|
71
|
+
w += [jnp.asarray(2.38*r[:len(i0)]/np.sqrt(n+1))]
|
|
72
|
+
b += [jnp.asarray(2.38*r[len(i0):]/np.sqrt(c+1))]
|
|
73
|
+
w = jnp.concatenate(w + b)
|
|
74
|
+
self.w = w
|
|
75
|
+
return w
|
|
76
|
+
|
|
77
|
+
def initx(self, x):
|
|
78
|
+
x = jnp.atleast_2d(x)
|
|
79
|
+
assert x.shape[1] == self.ni
|
|
80
|
+
ax, bx = utils.to01c(x)
|
|
81
|
+
self.f[0] = 'linear'
|
|
82
|
+
self.p[0] = (ax, bx)
|
|
83
|
+
return ax, bx
|
|
84
|
+
|
|
85
|
+
def inity(self, y):
|
|
86
|
+
y = jnp.atleast_2d(y)
|
|
87
|
+
assert y.shape[1] == self.no
|
|
88
|
+
ay, by = utils.from01c(y)
|
|
89
|
+
self.f[-1] = 'linear'
|
|
90
|
+
self.p[-1] = (ay, by)
|
|
91
|
+
return ay, by
|
|
92
|
+
|
|
93
|
+
def afunc(self, xi, i):
|
|
94
|
+
f = self.f[i]
|
|
95
|
+
p = self.p[i]
|
|
96
|
+
match f:
|
|
97
|
+
case 'identity':
|
|
98
|
+
return xi
|
|
99
|
+
case 'linear':
|
|
100
|
+
return self.linear(xi, *p)
|
|
101
|
+
case 'relu':
|
|
102
|
+
return jax.nn.relu(xi, *p)
|
|
103
|
+
case 'sigmoid':
|
|
104
|
+
return jax.nn.sigmoid(xi, *p)
|
|
105
|
+
case 'softplus':
|
|
106
|
+
return jax.nn.softplus(xi, *p)
|
|
107
|
+
case _:
|
|
108
|
+
return f(xi, *p)
|
|
109
|
+
|
|
110
|
+
def _prop(self, x):
|
|
111
|
+
y = self.afunc(x, 0)
|
|
112
|
+
Y = jnp.zeros(self.n).at[jnp.asarray(self.inps)].set(y)
|
|
113
|
+
for i, (s, (i0, i1), (j, k), (l, m)) in enumerate(self.wshp):
|
|
114
|
+
wi = self.w[j:k]
|
|
115
|
+
bi = self.w[l:m]
|
|
116
|
+
if self.sparse:
|
|
117
|
+
wi = BCOO((wi, jnp.asarray((i0, i1)).T), shape=s)
|
|
118
|
+
else:
|
|
119
|
+
wi = jnp.zeros(s).at[i0, i1].set(wi) # dense is default
|
|
120
|
+
yi = Y[:s[1]]
|
|
121
|
+
xi = wi@yi + bi
|
|
122
|
+
y = self.afunc(xi, i+1)
|
|
123
|
+
Y = Y.at[s[1]:s[1]+s[0]].set(y)
|
|
124
|
+
return Y[jnp.asarray(self.outs)]
|
|
125
|
+
|
|
126
|
+
@jax.jit
|
|
127
|
+
@partial(jax.vmap, in_axes=(None, 0, None))
|
|
128
|
+
def propd(self, x, spec):
|
|
129
|
+
f = self._prop
|
|
130
|
+
|
|
131
|
+
def deriv(f, x, v):
|
|
132
|
+
return jax.jvp(f, [x], [v])[1]
|
|
133
|
+
|
|
134
|
+
def one_hot(i):
|
|
135
|
+
return jnp.zeros(len(x)).at[i].set(1)
|
|
136
|
+
|
|
137
|
+
df = []
|
|
138
|
+
dy = []
|
|
139
|
+
for i, a, in enumerate(spec):
|
|
140
|
+
df += [[f]]
|
|
141
|
+
for j, n in enumerate(a[1:]):
|
|
142
|
+
def dfij(x, i=i, j=j, n=n):
|
|
143
|
+
vn = one_hot(n)
|
|
144
|
+
return deriv(df[i][j], x, vn)
|
|
145
|
+
df[i] += [dfij]
|
|
146
|
+
k = jnp.atleast_1d(jnp.asarray(a[0]))
|
|
147
|
+
dy += [df[i][-1](x)[k]]
|
|
148
|
+
return jnp.concatenate(dy)
|
|
149
|
+
|
|
150
|
+
@jax.jit
|
|
151
|
+
def propx(self, xspec):
|
|
152
|
+
return tuple(self.propd(a[0], [a[1:]]).ravel() for a in xspec)
|
|
153
|
+
|
|
154
|
+
@jax.jit
|
|
155
|
+
@partial(jax.vmap, in_axes=(None, 0))
|
|
156
|
+
def __call__(self, x):
|
|
157
|
+
return self._prop(x)
|
|
158
|
+
|
|
159
|
+
# @jax.jit
|
|
160
|
+
# def __call__(self, x, spec=None):
|
|
161
|
+
# spec = self.codes.get(spec, spec)
|
|
162
|
+
# return self._propd(x, spec)
|
|
163
|
+
|
|
164
|
+
@jax.jit
|
|
165
|
+
def loss(self, w, x, t):
|
|
166
|
+
self.w = jnp.asarray(w)
|
|
167
|
+
return self.mse(self(x), t)
|
|
168
|
+
|
|
169
|
+
@jax.jit
|
|
170
|
+
@partial(jax.grad, argnums=1)
|
|
171
|
+
def lossg(self, w, *args):
|
|
172
|
+
return self.loss(w, *args)
|
|
173
|
+
|
|
174
|
+
@jax.jit
|
|
175
|
+
@partial(jax.value_and_grad, argnums=1)
|
|
176
|
+
def lossvg(self, w, *args):
|
|
177
|
+
return self.loss(w, *args)
|
|
178
|
+
|
|
179
|
+
@jax.jit
|
|
180
|
+
@partial(jax.jacobian, argnums=1)
|
|
181
|
+
def eqcj(self, w, *args):
|
|
182
|
+
return self.eqc(w, *args)
|
|
183
|
+
|
|
184
|
+
@jax.jit
|
|
185
|
+
@partial(jax.jacobian, argnums=1)
|
|
186
|
+
def ineqcj(self, w, *args):
|
|
187
|
+
return self.ineqc(w, *args)
|
|
188
|
+
|
|
189
|
+
# @jax.jit
|
|
190
|
+
# def loss(self, w, *args):
|
|
191
|
+
# self.w = jnp.asarray(w)
|
|
192
|
+
# y = self.eqs(w, *args)
|
|
193
|
+
# return sum(mse(yi) for yi in y)
|
|
194
|
+
|
|
195
|
+
# This is concrete represntation of eqs
|
|
196
|
+
# @jax.jit
|
|
197
|
+
# def eqs(self, w, x, t):
|
|
198
|
+
# self.w = jnp.asarray(w)
|
|
199
|
+
# return [self(x) - t]
|
|
200
|
+
|
|
201
|
+
# Unfortunately jax's BFGS most of the time fails because of 3 - zoom error
|
|
202
|
+
# def train(self, *args, maxiter=None, tol=None, method='BFGS', **kwargs):
|
|
203
|
+
# from jax.scipy.optimize import minimize
|
|
204
|
+
# res = minimize(self.loss, self.w, args,
|
|
205
|
+
# tol=tol,
|
|
206
|
+
# options={'maxiter': maxiter},
|
|
207
|
+
# method=method,
|
|
208
|
+
# **kwargs)
|
|
209
|
+
# self.w = jnp.asarray(res.x)
|
|
210
|
+
# return res
|
|
211
|
+
|
|
212
|
+
# We use scipy.optimize.minimize here. It has many advantages,
|
|
213
|
+
# but cannot be jitted...
|
|
214
|
+
def train(self, *args, tol=1e-5, maxiter=None, disp=False, **kwargs):
|
|
215
|
+
options = kwargs.pop('options', {})
|
|
216
|
+
options.setdefault('maxiter', maxiter)
|
|
217
|
+
options.setdefault('disp', disp)
|
|
218
|
+
cstr = []
|
|
219
|
+
if hasattr(self, 'eqc'):
|
|
220
|
+
cstr += [{'type': 'eq',
|
|
221
|
+
'fun': self.eqc,
|
|
222
|
+
'jac': self.eqcj,
|
|
223
|
+
'args': args}]
|
|
224
|
+
if hasattr(self, 'ineqc'):
|
|
225
|
+
cstr += [{'type': 'ineq',
|
|
226
|
+
'fun': self.ineqc,
|
|
227
|
+
'jac': self.ineqcj,
|
|
228
|
+
'args': args}]
|
|
229
|
+
method = 'BFGS' if len(cstr) == 0 else 'SLSQP'
|
|
230
|
+
if method == 'SLSQP':
|
|
231
|
+
options['ftol'] = tol # it seems this is set here
|
|
232
|
+
res = minimize(self.loss, self.w, args, jac=self.lossg,
|
|
233
|
+
method=method,
|
|
234
|
+
tol=tol,
|
|
235
|
+
constraints=cstr,
|
|
236
|
+
options=options,
|
|
237
|
+
**kwargs)
|
|
238
|
+
self.w = jnp.asarray(res.x)
|
|
239
|
+
return res
|
|
240
|
+
|
|
241
|
+
@staticmethod
|
|
242
|
+
def linear(x, a=1., b=0.):
|
|
243
|
+
return a*x + b
|
|
244
|
+
|
|
245
|
+
@staticmethod
|
|
246
|
+
def mse(y, t=0.):
|
|
247
|
+
y = jnp.asarray(y).ravel()
|
|
248
|
+
t = jnp.asarray(t).ravel()
|
|
249
|
+
return jnp.mean((y - t)**2)
|
|
250
|
+
|
|
251
|
+
@staticmethod
|
|
252
|
+
def rmse(y, t=0.):
|
|
253
|
+
return jnp.sqrt(spinn.mse(y, t))
|
|
254
|
+
|
|
255
|
+
def _tree_flatten(self):
|
|
256
|
+
children = (self.w, self.p) # dynamic values
|
|
257
|
+
aux_data = {'nlst': self.nlst,
|
|
258
|
+
'n': self.n,
|
|
259
|
+
'inps': self.inps,
|
|
260
|
+
'outs': self.outs,
|
|
261
|
+
'ni': self.ni,
|
|
262
|
+
'no': self.no,
|
|
263
|
+
'nw': self.nw,
|
|
264
|
+
'nb': self.nb,
|
|
265
|
+
'wshp': self.wshp,
|
|
266
|
+
'f': self.f,
|
|
267
|
+
'sparse': self.sparse} # static values
|
|
268
|
+
return (children, aux_data)
|
|
269
|
+
|
|
270
|
+
@classmethod
|
|
271
|
+
def _tree_unflatten(cls, aux_data, children):
|
|
272
|
+
obj = cls(None)
|
|
273
|
+
obj.w = children[0]
|
|
274
|
+
obj.p = children[1]
|
|
275
|
+
obj.__dict__.update(aux_data)
|
|
276
|
+
return obj
|
|
277
|
+
|
|
278
|
+
@classmethod
|
|
279
|
+
def _register_pytree(cls):
|
|
280
|
+
# This method is called at __init__ and allows for
|
|
281
|
+
# automatic pytree registration of subclasses.
|
|
282
|
+
# Function 'register_pytree_node' returns ValueError when registering
|
|
283
|
+
# the same class many times, so there is the 'if' below. Anyway,
|
|
284
|
+
# in case of troubles try to do this explicitly!
|
|
285
|
+
if cls not in jax._src.tree_util._registry:
|
|
286
|
+
jax.tree_util.register_pytree_node(cls,
|
|
287
|
+
cls._tree_flatten,
|
|
288
|
+
cls._tree_unflatten)
|
|
289
|
+
|
|
290
|
+
|
|
291
|
+
if __name__ == "__main__":
|
|
292
|
+
import os
|
|
293
|
+
import subprocess
|
|
294
|
+
fname = os.path.dirname(os.path.realpath(__file__))
|
|
295
|
+
subprocess.call(['pytest', fname])
|
|
@@ -0,0 +1,166 @@
|
|
|
1
|
+
import itertools
|
|
2
|
+
import functools
|
|
3
|
+
import random
|
|
4
|
+
import jax
|
|
5
|
+
import jax.numpy as jnp
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def prod(s1, s2):
|
|
9
|
+
return list(itertools.product(s1, s2))
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def acc(s):
|
|
13
|
+
return list(itertools.accumulate(s))
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def pairwise(s):
|
|
17
|
+
return list(itertools.pairwise(s))
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def pacc(s):
|
|
21
|
+
return list(itertools.pairwise(itertools.accumulate(s)))
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def to01c(x):
|
|
25
|
+
assert len(x.shape) == 2
|
|
26
|
+
assert x.shape[0] >= 2
|
|
27
|
+
|
|
28
|
+
def _to01(lim):
|
|
29
|
+
a, b = lim.T
|
|
30
|
+
c0 = 1./(b-a)
|
|
31
|
+
c1 = -a * c0
|
|
32
|
+
return c0, c1
|
|
33
|
+
|
|
34
|
+
xlim = jnp.asarray([(i.min(), i.max()) for i in x.T])
|
|
35
|
+
ax, bx = _to01(xlim)
|
|
36
|
+
return ax, bx
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def from01c(x):
|
|
40
|
+
def _from01(lim):
|
|
41
|
+
c, d = lim.T
|
|
42
|
+
c0 = d - c
|
|
43
|
+
c1 = c
|
|
44
|
+
return c0, c1
|
|
45
|
+
|
|
46
|
+
if isinstance(x, int):
|
|
47
|
+
xlim = jnp.asarray([(0., 1.)]*x)
|
|
48
|
+
else:
|
|
49
|
+
xlim = jnp.asarray([(i.min(), i.max()) for i in x.T])
|
|
50
|
+
|
|
51
|
+
ax, bx = _from01(xlim)
|
|
52
|
+
return ax, bx
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def unroll_fp(f, p, arch):
|
|
56
|
+
nl = len(arch)
|
|
57
|
+
if f is None:
|
|
58
|
+
f = ['identity'] + ['softplus']*(nl-2) + ['identity']
|
|
59
|
+
elif isinstance(f, str) or callable(f):
|
|
60
|
+
f = ['identity'] + [f]*(nl-2) + ['identity']
|
|
61
|
+
elif isinstance(f, list):
|
|
62
|
+
if len(f) == 3:
|
|
63
|
+
f = [f[0]] + [f[1]]*(nl-2) + [f[2]]
|
|
64
|
+
elif len(f) == len(arch):
|
|
65
|
+
f = f
|
|
66
|
+
else:
|
|
67
|
+
raise ValueError(f"{len(f)} funcs given for {len(arch)} layers.")
|
|
68
|
+
else:
|
|
69
|
+
raise ValueError("'f' have to be None, str, callable or a list.")
|
|
70
|
+
if p is None:
|
|
71
|
+
p = [()]*nl
|
|
72
|
+
elif isinstance(p, tuple):
|
|
73
|
+
p = [()] + [p]*(nl-2) + [()]
|
|
74
|
+
elif isinstance(p, list):
|
|
75
|
+
if len(p) == 3:
|
|
76
|
+
p = [p[0]] + [p[1]]*(nl-2) + [p[2]]
|
|
77
|
+
elif len(p) == len(arch):
|
|
78
|
+
p = p
|
|
79
|
+
else:
|
|
80
|
+
raise ValueError(f"{len(p)} params given for {len(arch)} layers")
|
|
81
|
+
else:
|
|
82
|
+
raise ValueError("'p' have to be None, tuple or a list.")
|
|
83
|
+
return f, p
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def meshgrid(*arrs):
|
|
87
|
+
arrs = tuple(reversed(arrs))
|
|
88
|
+
lens = tuple(map(len, arrs))
|
|
89
|
+
dim = len(arrs)
|
|
90
|
+
sz = 1
|
|
91
|
+
for s in lens:
|
|
92
|
+
sz *= s
|
|
93
|
+
ans = []
|
|
94
|
+
for i, arr in enumerate(arrs):
|
|
95
|
+
slc = [1]*dim
|
|
96
|
+
slc[i] = lens[i]
|
|
97
|
+
arr2 = jnp.asarray(arr).reshape(slc)
|
|
98
|
+
for j, sz in enumerate(lens):
|
|
99
|
+
if j != i:
|
|
100
|
+
arr2 = arr2.repeat(sz, axis=j)
|
|
101
|
+
ans.append(arr2)
|
|
102
|
+
return tuple(ans)
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def hcube(dims, n=11, origin=None):
|
|
106
|
+
dims = (1.,)*dims if isinstance(dims, int) else dims
|
|
107
|
+
dims = (dims,) if isinstance(dims, float) else dims
|
|
108
|
+
d = len(dims)
|
|
109
|
+
n = (n,)*d if isinstance(n, int) else n
|
|
110
|
+
origin = (0,)*d if origin is None else origin
|
|
111
|
+
xyz = [jnp.linspace(origin[i], dims[i]+origin[i], n[i]) for i in range(d)]
|
|
112
|
+
XYZ = meshgrid(*xyz)
|
|
113
|
+
return jnp.vstack(tuple(map(jnp.ravel, XYZ))).T
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def left(pts):
|
|
117
|
+
x = pts[:, 0]
|
|
118
|
+
return pts[jnp.isclose(x, min(x))]
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def right(pts):
|
|
122
|
+
x = pts[:, 0]
|
|
123
|
+
return pts[jnp.isclose(x, max(x))]
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def bottom(pts):
|
|
127
|
+
x = pts[:, 1]
|
|
128
|
+
return pts[jnp.isclose(x, min(x))]
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
def top(pts):
|
|
132
|
+
x = pts[:, 1]
|
|
133
|
+
return pts[jnp.isclose(x, max(x))]
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
def uniform(low=0., high=1., shape=(), key=None):
|
|
137
|
+
shape = (shape,) if isinstance(shape, int) else shape
|
|
138
|
+
key = random.randint(0, 10000000000) if key is None else key
|
|
139
|
+
key = jax.random.PRNGKey(key) if isinstance(key, int) else key
|
|
140
|
+
return jax.random.uniform(key, shape=shape, minval=low, maxval=high)
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def randomweights(arch, key=None):
|
|
144
|
+
W = []
|
|
145
|
+
wshp = list(itertools.pairwise(arch))
|
|
146
|
+
bshp = arch[1:]
|
|
147
|
+
for i, n in enumerate(arch[:-1]):
|
|
148
|
+
bound = 2.38/jnp.sqrt(n+1)
|
|
149
|
+
w = uniform(-bound, bound, shape=wshp[i])
|
|
150
|
+
b = uniform(-bound, bound, shape=bshp[i])
|
|
151
|
+
W += [[w, b]]
|
|
152
|
+
return W
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
@jax.jit
|
|
156
|
+
@functools.partial(jax.vmap, in_axes=(None, 0, None))
|
|
157
|
+
def simpleprop(w, x, c):
|
|
158
|
+
ax, bx, ay, by = c
|
|
159
|
+
y = ax*x + bx
|
|
160
|
+
for wi, bi in w[:-1]:
|
|
161
|
+
yi = y.dot(wi) + bi
|
|
162
|
+
y = jax.nn.softplus(yi)
|
|
163
|
+
wi, bi = w[-1]
|
|
164
|
+
y = y.dot(wi) + bi
|
|
165
|
+
y = ay*y + by
|
|
166
|
+
return y
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: spinn
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Simple Physics Informed Neural Networks
|
|
5
|
+
Home-page: https://github.com/mrkwjc/spinn
|
|
6
|
+
Author: Marek Wojciechowski
|
|
7
|
+
Author-email: mwojc@p.lodz.pl
|
|
8
|
+
License: LGPL-2.1
|
|
9
|
+
Keywords: ann,pinn
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Operating System :: OS Independent
|
|
13
|
+
Requires-Python: >=3.6
|
|
14
|
+
Description-Content-Type: text/markdown
|
|
15
|
+
License-File: LICENSE
|
|
16
|
+
Requires-Dist: jax
|
|
17
|
+
Requires-Dist: jaxlib
|
|
18
|
+
Requires-Dist: numpy
|
|
19
|
+
Requires-Dist: scipy
|
|
20
|
+
Requires-Dist: networkx
|
|
21
|
+
|
|
22
|
+
# spinn - Simple Physics Informed Neural Networks
|
|
23
|
+
Makes PINNs easy (with jax), see examples...
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
setup.py
|
|
5
|
+
spinn/__init__.py
|
|
6
|
+
spinn/graphs.py
|
|
7
|
+
spinn/trees.py
|
|
8
|
+
spinn/utils.py
|
|
9
|
+
spinn.egg-info/PKG-INFO
|
|
10
|
+
spinn.egg-info/SOURCES.txt
|
|
11
|
+
spinn.egg-info/dependency_links.txt
|
|
12
|
+
spinn.egg-info/requires.txt
|
|
13
|
+
spinn.egg-info/top_level.txt
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
spinn
|