Metadata-Version: 2.4
Name: spinn
Version: 0.3.0
Summary: Simple Physics Informed Neural Networks
Home-page: https://github.com/mrkwjc/spinn
Author: Marek Wojciechowski
Author-email: mwojc@p.lodz.pl
License: MIT
Keywords: ann,pinn
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax
Requires-Dist: jaxlib
Requires-Dist: numpy
Requires-Dist: scipy
Requires-Dist: networkx
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: keywords
Dynamic: license
Dynamic: license-file
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# spinn - Simple Physics Informed Neural Networks
Makes PINNs easy (with jax), see examples...

## Usage

Subclass `spinn` and define `terms`: the loss terms of the problem, with the
equation (PDE residual) first and the conditions after it. `propx` returns
the network outputs and their derivatives at given points: `[x, k, i, j, ...]`
is output `k` at the points `x`, differentiated with respect to inputs
`i`, `j`, ...

```python
import jax
import jax.numpy as jnp
from spinn import spinn, hcube, left, right, bottom, mlgraph


class heat(spinn):
    """dF/dt = a d2F/dx2, F = 0 at x = 0 and x = 1, F = sin(pi x) at t = 0"""

    @jax.jit
    def terms(self, w, x, a, f0, bx):
        self.w = jnp.asarray(w)
        left, right, bottom = bx
        dFt, dFxx = self.propx([[x, 0, 1], [x, 0, 0, 0]])
        Fl, Fr, Fb = self.propx([[left, 0], [right, 0], [bottom, 0]])
        return (self.mse(dFt - a*dFxx),             # equation
                self.mse(Fl) + self.mse(Fr),        # boundary conditions
                self.mse(Fb, f0))                   # initial condition


x = hcube((1., 0.5), n=21)                  # grid of points (x, t)
bx = left(x), right(x), bottom(x)
f0 = jnp.sin(jnp.pi*bx[2][:, 0])

P = heat(mlgraph((2, 16, 16, 1)), initx=x)  # 2 inputs, 2 hidden layers
res = P.train(x, 0.2, f0, bx)
F = P(x)                                    # solution at the points
```

`train` minimizes the weighted sum of the terms with scipy's L-BFGS-B. The
weights are balanced automatically (`balance='curv'` by default, see
`spinn.balance` for the other rules, e.g. `balance='curvsens'`), or given
explicitly as `balance=[1., 10., 10.]`. The final weights are returned in
`res.lmbd` and the unweighted terms in `res.terms`.

Activations of the hidden layers are set with `f`: a name from `jax.nn` or
`jax.numpy` (`f='tanh'`, `f='sin'`, `f='gelu'`; softplus by default), a
function, or one per layer. Networks are saved with `P.dump(fname)` and
restored with `heat.load(fname)`; see `help(spinn)` for details.

More problems (beams, Laplace and heat equations) are in `examples/`.
