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 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
@@ -0,0 +1,2 @@
1
+ # spinn - Simple Physics Informed Neural Networks
2
+ Makes PINNs easy (with jax), see examples...
@@ -0,0 +1,3 @@
1
+ [build-system]
2
+ requires=["setuptools"]
3
+ build-backend = "setuptools.build_meta"
spinn-0.1.0/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
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,6 @@
1
+ import jax
2
+ jax.config.update('jax_default_matmul_precision', 'high') # 'bfloat16_3x'
3
+
4
+ from .trees import spinn
5
+ from .utils import hcube, left, right, bottom, top
6
+ from .graphs import mlgraph
@@ -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,5 @@
1
+ jax
2
+ jaxlib
3
+ numpy
4
+ scipy
5
+ networkx
@@ -0,0 +1 @@
1
+ spinn